@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 237 lines 7.84 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import functools
5 import hashlib
6 import pickle
7 import uuid
8 from collections.abc import Callable, Iterable
9
10 import numpy as np
11 import torch
12 from PIL import Image
13
14 from vllm.config.multimodal import MMHasherAlgorithm
15 from vllm.logger import init_logger
16
17 from .media import MediaWithBytes
18
19 logger = init_logger(__name__)
20
21 # Framing for the digest input. The hash is built by feeding a stream of byte
22 # chunks to `hasher.update`, so the stream must be uniquely decodable: without
23 # an explicit length in front of every chunk and a tag in front of every
24 # container, distinct inputs serialize to identical bytes and share a cache
25 # entry. See `test_hash_collision_*` in tests/multimodal/test_hasher.py.
26 _LENGTH_BYTES = 8
27 _TAG_NONE = b"\x00"
28 _TAG_SEQUENCE = b"\x01"
29 _TAG_MAPPING = b"\x02"
30 _TAG_LEAF = b"\x03"
31
32
33 def _encode_length(value: int) -> bytes:
34 return value.to_bytes(_LENGTH_BYTES, "little")
35
36
37 def _framed(chunk: bytes | memoryview) -> Iterable[bytes | memoryview]:
38 """Yield *chunk* preceded by its size in bytes."""
39 size = chunk.nbytes if isinstance(chunk, memoryview) else len(chunk)
40 yield _encode_length(size)
41 yield chunk
42
43
44 @functools.lru_cache(maxsize=3)
45 def _get_hasher_factory(
46 algorithm: MMHasherAlgorithm,
47 ) -> Callable[[], "hashlib._Hash"]:
48 """Get the hasher factory based on the configured algorithm.
49
50 Args:
51 algorithm: Hash algorithm name (blake3, sha256, or sha512)
52
53 Returns a callable that creates a new hasher instance.
54 Supports blake3 (default), sha256, and sha512 for FIPS compliance.
55
56 See: https://github.com/vllm-project/vllm/issues/18334
57
58 """
59 if algorithm == "blake3":
60 from blake3 import blake3
61
62 return blake3
63 elif algorithm == "sha256":
64 return hashlib.sha256
65 elif algorithm == "sha512":
66 return hashlib.sha512
67 else:
68 # This should never happen due to config validation
69 raise ValueError(f"Unsupported hash algorithm: {algorithm}")
70
71
72 def _get_image_id_bytes(image: Image.Image) -> bytes | None:
73 try:
74 exif = image.getexif()
75 image_id = exif.get(Image.ExifTags.Base.ImageID)
76 if isinstance(image_id, uuid.UUID):
77 return image_id.bytes
78 except Exception:
79 # Tolerate malformed EXIF metadata (e.g. invalid TIFF header)
80 # and fall back to serializing raw image data or bytes.
81 pass
82 return None
83
84
85 class MultiModalHasher:
86 """Derives multi-modal cache keys.
87
88 Every method here yields *framed* chunks: each chunk is preceded by its
89 length and each container by its kind, so that the concatenation fed to the
90 digest is uniquely decodable and distinct inputs cannot share a key.
91 """
92
93 @classmethod
94 def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]:
95 # Simple cases
96 if isinstance(obj, (bytes, memoryview)):
97 return _framed(obj)
98 if isinstance(obj, str):
99 return _framed(obj.encode("utf-8"))
100 if isinstance(obj, (int, float)):
101 return _framed(np.array(obj).tobytes())
102
103 if isinstance(obj, Image.Image):
104 image_id = _get_image_id_bytes(obj)
105 if image_id is not None:
106 return _framed(image_id)
107
108 data = {"mode": obj.mode, "data": np.asarray(obj)}
109 palette = obj.palette
110 if palette is not None:
111 data["palette"] = palette.palette
112 if palette.rawmode is not None:
113 data["palette_rawmode"] = palette.rawmode
114
115 return cls.iter_item_to_bytes("image", data)
116
117 if isinstance(obj, MediaWithBytes) and isinstance(obj.media, Image.Image):
118 image_id = _get_image_id_bytes(obj.media)
119 if image_id is not None:
120 return _framed(image_id)
121
122 if obj.io_config:
123 return cls.iter_item_to_bytes(
124 "image",
125 {"io_config": obj.io_config, "data": obj.original_bytes},
126 )
127 return cls.iter_item_to_bytes("image", obj.original_bytes)
128
129 if isinstance(obj, MediaWithBytes) and isinstance(
130 obj.media, (np.ndarray, torch.Tensor)
131 ):
132 frames = obj.media
133 # Both np.ndarray and torch.Tensor expose .nbytes.
134 if frames.nbytes < len(obj.original_bytes):
135 return cls.iter_item_to_bytes("video", frames)
136 return cls.iter_item_to_bytes("video", obj.original_bytes)
137
138 if isinstance(obj, torch.Tensor):
139 tensor_obj: torch.Tensor = obj.cpu()
140 tensor_dtype = tensor_obj.dtype
141 tensor_shape = tensor_obj.shape
142
143 # NumPy does not support bfloat16.
144 # Workaround: View the tensor as a contiguous 1D array of bytes
145 if tensor_dtype == torch.bfloat16:
146 tensor_obj = tensor_obj.contiguous()
147 tensor_obj = tensor_obj.view((tensor_obj.numel(),)).view(torch.uint8)
148
149 return cls.iter_item_to_bytes(
150 "tensor",
151 {
152 "original_dtype": str(tensor_dtype),
153 "original_shape": tuple(tensor_shape),
154 "data": tensor_obj.numpy(),
155 },
156 )
157
158 return cls.iter_item_to_bytes("tensor", tensor_obj.numpy())
159
160 if isinstance(obj, np.ndarray):
161 if obj.ndim == 0:
162 arr_data = obj.item()
163 elif obj.flags.c_contiguous:
164 # Not valid for 0-D arrays
165 arr_data = obj.view(np.uint8).data
166 else:
167 # If the array is non-contiguous, we need to copy it first
168 arr_data = obj.tobytes()
169
170 return cls.iter_item_to_bytes(
171 "ndarray",
172 {
173 "dtype": obj.dtype.str,
174 "shape": obj.shape,
175 "data": arr_data,
176 },
177 )
178
179 logger.warning(
180 "No serialization method found for %s. Falling back to pickle.", type(obj)
181 )
182
183 return _framed(pickle.dumps(obj))
184
185 @classmethod
186 def iter_item_to_bytes(
187 cls,
188 key: str,
189 obj: object,
190 ) -> Iterable[bytes | memoryview]:
191 """Yield the digest input for a single ``key``/``obj`` pair."""
192 yield from _framed(key.encode("utf-8"))
193 yield from cls.iter_value_to_bytes(obj)
194
195 @classmethod
196 def iter_value_to_bytes(
197 cls,
198 obj: object,
199 ) -> Iterable[bytes | memoryview]:
200 """Yield the digest input for a value, tagged by its container kind.
201
202 Containers carry their kind and length so that a nested structure can
203 never serialize to the same bytes as a differently shaped one (a list
204 and a mapping keyed by stringified indices, for example).
205 """
206 if obj is None:
207 yield _TAG_NONE
208 elif isinstance(obj, (list, tuple)):
209 yield _TAG_SEQUENCE
210 yield _encode_length(len(obj))
211 for elem in obj:
212 yield from cls.iter_value_to_bytes(elem)
213 elif isinstance(obj, dict):
214 yield _TAG_MAPPING
215 yield _encode_length(len(obj))
216 for k, v in obj.items():
217 yield from _framed(str(k).encode("utf-8"))
218 yield from cls.iter_value_to_bytes(v)
219 else:
220 yield _TAG_LEAF
221 yield from cls.serialize_item(obj)
222
223 @classmethod
224 def hash_kwargs(
225 cls,
226 algorithm: MMHasherAlgorithm,
227 /,
228 **kwargs: object,
229 ) -> str:
230 hasher_factory = _get_hasher_factory(algorithm)
231 hasher = hasher_factory()
232
233 for k, v in sorted(kwargs.items(), key=lambda kv: kv[0]):
234 for bytes_ in cls.iter_item_to_bytes(k, v):
235 hasher.update(bytes_)
236
237 return hasher.hexdigest()