@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 250 lines 7.83 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from typing import Any
5
6 import msgspec
7 import pytest
8
9 from vllm.distributed.kv_events import BlockRemoved, BlockStored
10
11 # Minimal ExternalBlockHash for testing (bytes are a valid ExternalBlockHash).
12 _FAKE_HASH: bytes = b"\xab" * 32
13
14
15 class _LegacyBlockStored(
16 msgspec.Struct,
17 omit_defaults=True, # type: ignore[call-arg]
18 gc=False, # type: ignore[call-arg]
19 tag="BlockStored", # type: ignore[call-arg]
20 ):
21 """BlockStored wire schema before locality was added."""
22
23 block_hashes: list[bytes]
24 parent_block_hash: bytes | None
25 token_ids: list[int]
26 block_size: int
27 lora_id: int | None
28 medium: str | None
29 lora_name: str | None
30 extra_keys: list[tuple[Any, ...] | None] | None = None
31 group_idx: int | None = None
32 kv_cache_spec_kind: str | None = None
33 kv_cache_spec_sliding_window: int | None = None
34
35
36 class _LegacyBlockRemoved(
37 msgspec.Struct,
38 omit_defaults=True, # type: ignore[call-arg]
39 gc=False, # type: ignore[call-arg]
40 tag="BlockRemoved", # type: ignore[call-arg]
41 ):
42 """BlockRemoved wire schema before locality was added."""
43
44 block_hashes: list[bytes]
45 medium: str | None
46 group_idx: int | None = None
47
48
49 class _PreSessionBlockStored(
50 msgspec.Struct,
51 omit_defaults=True, # type: ignore[call-arg]
52 gc=False, # type: ignore[call-arg]
53 tag="BlockStored", # type: ignore[call-arg]
54 ):
55 """BlockStored wire schema before session_id was added."""
56
57 block_hashes: list[bytes]
58 parent_block_hash: bytes | None
59 token_ids: list[int]
60 block_size: int
61 lora_id: int | None
62 medium: str | None
63 lora_name: str | None
64 extra_keys: list[tuple[Any, ...] | None] | None = None
65 group_idx: int | None = None
66 kv_cache_spec_kind: str | None = None
67 kv_cache_spec_sliding_window: int | None = None
68 locality: str | None = None
69
70
71 def _make_block_stored(
72 group_idx: int | None = None,
73 kv_cache_spec_sliding_window: int | None = None,
74 locality: str | None = None,
75 session_id: str | None = None,
76 ) -> BlockStored:
77 return BlockStored(
78 block_hashes=[_FAKE_HASH],
79 parent_block_hash=None,
80 token_ids=[1, 2, 3, 4],
81 block_size=4,
82 lora_id=None,
83 medium="GPU",
84 lora_name=None,
85 group_idx=group_idx,
86 kv_cache_spec_sliding_window=kv_cache_spec_sliding_window,
87 locality=locality,
88 session_id=session_id,
89 )
90
91
92 def _make_block_removed(
93 group_idx: int | None = None,
94 locality: str | None = None,
95 ) -> BlockRemoved:
96 return BlockRemoved(
97 block_hashes=[_FAKE_HASH],
98 medium="GPU",
99 group_idx=group_idx,
100 locality=locality,
101 )
102
103
104 def test_block_stored_default_group_idx_is_none():
105 """group_idx defaults to None when not provided."""
106 event = _make_block_stored()
107 assert event.group_idx is None
108
109
110 def test_block_removed_default_group_idx_is_none():
111 """group_idx defaults to None when not provided."""
112 event = _make_block_removed()
113 assert event.group_idx is None
114
115
116 @pytest.mark.parametrize("group_idx", [1, 2, 3])
117 def test_block_stored_hash_differs_by_group_idx(group_idx: int):
118 """BlockStored events that differ only in group_idx must hash differently."""
119 other_group_idx = group_idx + 1
120 event_a = _make_block_stored(group_idx=group_idx)
121 event_b = _make_block_stored(group_idx=other_group_idx)
122 assert hash(event_a) != hash(event_b)
123
124
125 def test_block_stored_hash_same_for_equal_group_idx():
126 """Two BlockStored events with identical fields produce the same hash."""
127 event_a = _make_block_stored(group_idx=1)
128 event_b = _make_block_stored(group_idx=1)
129 assert hash(event_a) == hash(event_b)
130
131
132 def test_block_stored_hash_differs_by_session_id():
133 event_a = _make_block_stored(session_id="session-a")
134 event_b = _make_block_stored(session_id="session-b")
135 assert hash(event_a) != hash(event_b)
136
137
138 @pytest.mark.parametrize("group_idx", [1, 2, 3])
139 def test_block_removed_hash_differs_by_group_idx(group_idx: int):
140 """BlockRemoved events that differ only in group_idx must hash differently."""
141 other_group_idx = group_idx + 1
142 event_a = _make_block_removed(group_idx=group_idx)
143 event_b = _make_block_removed(group_idx=other_group_idx)
144 assert hash(event_a) != hash(event_b)
145
146
147 def test_block_removed_hash_same_for_equal_group_idx():
148 """Two BlockRemoved events with identical fields produce the same hash."""
149 event_a = _make_block_removed(group_idx=1)
150 event_b = _make_block_removed(group_idx=1)
151 assert hash(event_a) == hash(event_b)
152
153
154 def test_block_stored_hash_differs_by_sliding_window():
155 event_a = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=128)
156 event_b = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=256)
157 assert hash(event_a) != hash(event_b)
158
159
160 @pytest.mark.parametrize(
161 ("event_a", "event_b"),
162 [
163 (
164 _make_block_stored(locality="LOCAL"),
165 _make_block_stored(locality="REMOTE"),
166 ),
167 (
168 _make_block_removed(locality="LOCAL"),
169 _make_block_removed(locality="REMOTE"),
170 ),
171 ],
172 )
173 def test_event_hash_differs_by_locality(
174 event_a: BlockStored | BlockRemoved,
175 event_b: BlockStored | BlockRemoved,
176 ):
177 assert hash(event_a) != hash(event_b)
178
179
180 def test_block_stored_locality_is_wire_compatible():
181 legacy = _LegacyBlockStored(
182 block_hashes=[_FAKE_HASH],
183 parent_block_hash=None,
184 token_ids=[1, 2, 3, 4],
185 block_size=4,
186 lora_id=None,
187 medium="GPU",
188 lora_name=None,
189 group_idx=2,
190 kv_cache_spec_sliding_window=128,
191 )
192 legacy_payload = msgspec.msgpack.encode(legacy)
193 assert (
194 msgspec.msgpack.encode(
195 _make_block_stored(
196 group_idx=2,
197 kv_cache_spec_sliding_window=128,
198 )
199 )
200 == legacy_payload
201 )
202 assert msgspec.msgpack.decode(legacy_payload, type=BlockStored).locality is None
203 new_payload = msgspec.msgpack.encode(_make_block_stored(locality="LOCAL"))
204 assert msgspec.msgpack.decode(new_payload)["locality"] == "LOCAL"
205 assert msgspec.msgpack.decode(new_payload, type=_LegacyBlockStored).medium == "GPU"
206
207
208 def test_block_removed_locality_is_wire_compatible():
209 legacy = _LegacyBlockRemoved(block_hashes=[_FAKE_HASH], medium="GPU")
210 legacy_payload = msgspec.msgpack.encode(legacy)
211 assert msgspec.msgpack.encode(_make_block_removed()) == legacy_payload
212 assert msgspec.msgpack.decode(legacy_payload, type=BlockRemoved).locality is None
213 new_payload = msgspec.msgpack.encode(_make_block_removed(locality="REMOTE"))
214 assert msgspec.msgpack.decode(new_payload)["locality"] == "REMOTE"
215 assert msgspec.msgpack.decode(new_payload, type=_LegacyBlockRemoved).medium == "GPU"
216
217
218 def test_block_stored_session_id_is_wire_compatible():
219 pre_session = _PreSessionBlockStored(
220 block_hashes=[_FAKE_HASH],
221 parent_block_hash=None,
222 token_ids=[1, 2, 3, 4],
223 block_size=4,
224 lora_id=None,
225 medium="GPU",
226 lora_name=None,
227 group_idx=2,
228 kv_cache_spec_sliding_window=128,
229 locality="LOCAL",
230 )
231 pre_session_payload = msgspec.msgpack.encode(pre_session)
232 assert (
233 msgspec.msgpack.encode(
234 _make_block_stored(
235 group_idx=2,
236 kv_cache_spec_sliding_window=128,
237 locality="LOCAL",
238 )
239 )
240 == pre_session_payload
241 )
242 assert (
243 msgspec.msgpack.decode(pre_session_payload, type=BlockStored).session_id is None
244 )
245
246 new_payload = msgspec.msgpack.encode(_make_block_stored(session_id="session-1"))
247 assert msgspec.msgpack.decode(new_payload)["session_id"] == "session-1"
248 assert (
249 msgspec.msgpack.decode(new_payload, type=_PreSessionBlockStored).medium == "GPU"
250 )