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, 12:15 UTC
[Security] Fix image hash collision (#17378)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
May 2, 2025 at 20:36 UTC
99404f53c72965b41558aceb1bc2380875f5d848
4 files changed
+83
-10
tests/multimodal/assets/image1.png
Binary files /dev/null and b/tests/multimodal/assets/image1.png differ
tests/multimodal/assets/image2.png
Binary files /dev/null and b/tests/multimodal/assets/image2.png differ
tests/multimodal/test_hasher.py
new
+61
@@ -0,0 +1,61 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from pathlib import Path
3
+
4
+import numpy as np
5
+import pytest
6
+import torch
7
+from PIL import Image, ImageDraw
8
+
9
+from vllm.multimodal.hasher import MultiModalHasher
10
+
11
+ASSETS_DIR = Path(__file__).parent / "assets"
12
+assert ASSETS_DIR.exists()
13
+
14
+
15
+# NOTE: Images that are the same visually are allowed to have the same hash
16
+@pytest.mark.parametrize("mode_pair", [("1", "L"), ("RGBA", "CMYK")])
17
+def test_hash_collision_image_mode(mode_pair):
18
+ mode1, mode2 = mode_pair
19
+ image1 = Image.new(mode1, size=(10, 10), color=1)
20
+ image2 = Image.new(mode2, size=(10, 10), color=1)
21
+
22
+ hasher = MultiModalHasher
23
+ assert hasher.hash_kwargs(image=image1) != hasher.hash_kwargs(image=image2)
24
+
25
+
26
+def test_hash_collision_image_palette():
27
+ # These images differ only in Image.palette._palette
28
+ image1 = Image.open(ASSETS_DIR / "image1.png")
29
+ image2 = Image.open(ASSETS_DIR / "image2.png")
30
+
31
+ hasher = MultiModalHasher
32
+ assert hasher.hash_kwargs(image=image1) != hasher.hash_kwargs(image=image2)
33
+
34
+
35
+def test_hash_collision_image_transpose():
36
+ image1 = Image.new("1", size=(10, 20))
37
+ ImageDraw.Draw(image1).line([(0, 0), (10, 0)])
38
+
39
+ image2 = Image.new("1", size=(20, 10))
40
+ ImageDraw.Draw(image2).line([(0, 0), (0, 10)])
41
+
42
+ hasher = MultiModalHasher
43
+ assert hasher.hash_kwargs(image=image1) != hasher.hash_kwargs(image=image2)
44
+
45
+
46
+def test_hash_collision_tensor_shape():
47
+ # The hash should be different though the data is the same when flattened
48
+ arr1 = torch.zeros((5, 10, 20, 3))
49
+ arr2 = torch.zeros((10, 20, 5, 3))
50
+
51
+ hasher = MultiModalHasher
52
+ assert hasher.hash_kwargs(data=arr1) != hasher.hash_kwargs(data=arr2)
53
+
54
+
55
+def test_hash_collision_array_shape():
56
+ # The hash should be different though the data is the same when flattened
57
+ arr1 = np.zeros((5, 10, 20, 3))
58
+ arr2 = np.zeros((10, 20, 5, 3))
59
+
60
+ hasher = MultiModalHasher
61
+ assert hasher.hash_kwargs(data=arr1) != hasher.hash_kwargs(data=arr2)
vllm/multimodal/hasher.py
+22
-10
@@ -31,16 +31,20 @@ class MultiModalHasher:
31
return obj.encode("utf-8")
32
if isinstance(obj, bytes):
33
return obj
34
- if isinstance(obj, Image.Image):
35
- return obj.tobytes()
34
+ if isinstance(obj, (int, float)):
35
+ return np.array(obj).tobytes()
36
37
- # Convertible to NumPy arrays
37
+ if isinstance(obj, Image.Image):
38
+ return cls.item_to_bytes("image", np.array(obj.convert("RGBA")))
39
if isinstance(obj, torch.Tensor):
39
- obj = obj.numpy()
40
- if isinstance(obj, (int, float)):
41
- obj = np.array(obj)
40
+ return cls.item_to_bytes("tensor", obj.numpy())
41
if isinstance(obj, np.ndarray):
43
- return obj.tobytes()
42
+ return cls.item_to_bytes(
43
+ "ndarray", {
44
+ "dtype": obj.dtype.str,
45
+ "shape": obj.shape,
46
+ "data": obj.data.tobytes(),
47
+ })
48
49
logger.warning(
50
"No serialization method found for %s. "
@@ -53,14 +57,22 @@ class MultiModalHasher:
57
cls,
58
key: str,
59
obj: object,
60
+ ) -> bytes:
61
+ return b''.join(kb + vb for kb, vb in cls.iter_item_to_bytes(key, obj))
62
+
63
+ @classmethod
64
+ def iter_item_to_bytes(
65
+ cls,
66
+ key: str,
67
+ obj: object,
68
) -> Iterable[tuple[bytes, bytes]]:
69
# Recursive cases
70
if isinstance(obj, (list, tuple)):
71
for i, elem in enumerate(obj):
60
- yield from cls.item_to_bytes(f"{key}.{i}", elem)
72
+ yield from cls.iter_item_to_bytes(f"{key}.{i}", elem)
73
elif isinstance(obj, dict):
74
for k, v in obj.items():
63
- yield from cls.item_to_bytes(f"{key}.{k}", v)
75
+ yield from cls.iter_item_to_bytes(f"{key}.{k}", v)
76
else:
77
key_bytes = cls.serialize_item(key)
78
value_bytes = cls.serialize_item(obj)
@@ -71,7 +83,7 @@ class MultiModalHasher:
83
hasher = blake3()
84
85
for k, v in kwargs.items():
74
- for k_bytes, v_bytes in cls.item_to_bytes(k, v):
86
+ for k_bytes, v_bytes in cls.iter_item_to_bytes(k, v):
87
hasher.update(k_bytes)
88
hasher.update(v_bytes)
89