| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import torch |
| 5 | |
| 6 | from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange |
| 7 | from vllm.v1.core.encoder_cache_manager import ( |
| 8 | EncoderCacheManager, |
| 9 | EncoderDecoderCacheManager, |
| 10 | ) |
| 11 | |
| 12 | pytestmark = pytest.mark.cpu_test |
| 13 | |
| 14 | |
| 15 | # ------------------ Mock Classes ------------------ # |
| 16 | class MockRequest: |
| 17 | def __init__(self, request_id, mm_hashes, token_counts): |
| 18 | self.request_id = request_id |
| 19 | self._token_counts = token_counts |
| 20 | self.mm_features = [] |
| 21 | for i, mm_hash in enumerate(mm_hashes): |
| 22 | feature = MultiModalFeatureSpec( |
| 23 | data=None, |
| 24 | modality="image", |
| 25 | identifier=mm_hash, |
| 26 | mm_position=PlaceholderRange(offset=0, length=self._token_counts[i]), |
| 27 | ) |
| 28 | self.mm_features.append(feature) |
| 29 | |
| 30 | def get_num_encoder_embeds(self, input_id: int) -> int: |
| 31 | return self._token_counts[input_id] |
| 32 | |
| 33 | |
| 34 | # ------------------ Unit Tests ------------------ # |
| 35 | def test_basic_allocate_and_reuse(): |
| 36 | cache = EncoderCacheManager(cache_size=10) |
| 37 | req = MockRequest("r1", ["imgA"], [4]) |
| 38 | |
| 39 | assert not cache.check_and_update_cache(req, 0) |
| 40 | assert cache.can_allocate(req, 0, int(1e9), 0) |
| 41 | |
| 42 | cache.allocate(req, 0) |
| 43 | |
| 44 | assert cache.check_and_update_cache(req, 0) |
| 45 | assert "r1" in cache.cached["imgA"] |
| 46 | assert cache.request_cached_ids["r1"] == {0} |
| 47 | assert cache.num_free_slots == 6 |
| 48 | |
| 49 | # Free twice to bring refcount to 0. |
| 50 | cache.free_encoder_input(req, 0) |
| 51 | cache.free_encoder_input(req, 0) |
| 52 | |
| 53 | assert not cache.cached["imgA"] |
| 54 | assert "r1" not in cache.request_cached_ids |
| 55 | assert "imgA" in cache.freeable |
| 56 | assert cache.num_freeable_slots == 10 |
| 57 | assert cache.num_free_slots == 6 |
| 58 | |
| 59 | |
| 60 | def test_freeing_decreases_refcount_and_moves_to_freeable(): |
| 61 | manager = EncoderCacheManager(cache_size=10) |
| 62 | req = MockRequest("req2", ["img3"], [5]) |
| 63 | |
| 64 | assert manager.can_allocate(req, 0, int(1e9), 0) |
| 65 | manager.allocate(req, 0) |
| 66 | |
| 67 | assert len(manager.cached["img3"]) == 1 |
| 68 | assert manager.request_cached_ids["req2"] == {0} |
| 69 | |
| 70 | manager.free_encoder_input(req, 0) |
| 71 | |
| 72 | assert not manager.cached["img3"] |
| 73 | assert "req2" not in manager.request_cached_ids |
| 74 | assert "img3" in manager.freeable |
| 75 | assert manager.num_freeable_slots == 10 |
| 76 | |
| 77 | |
| 78 | def test_free_request_frees_all_inputs(): |
| 79 | manager = EncoderCacheManager(cache_size=10) |
| 80 | req = MockRequest("req3", ["a", "b"], [2, 3]) |
| 81 | |
| 82 | assert manager.can_allocate(req, 0, int(1e9), 0) |
| 83 | manager.allocate(req, 0) |
| 84 | |
| 85 | assert manager.can_allocate(req, 1, int(1e9), 0) |
| 86 | manager.allocate(req, 1) |
| 87 | |
| 88 | assert len(manager.cached["a"]) == 1 |
| 89 | assert len(manager.cached["b"]) == 1 |
| 90 | assert manager.request_cached_ids["req3"] == {0, 1} |
| 91 | |
| 92 | manager.free(req) |
| 93 | |
| 94 | assert not manager.cached["a"] |
| 95 | assert not manager.cached["b"] |
| 96 | assert "req3" not in manager.request_cached_ids |
| 97 | assert "a" in manager.freeable |
| 98 | assert "b" in manager.freeable |
| 99 | assert manager.num_freeable_slots == 10 |
| 100 | |
| 101 | |
| 102 | def test_eviction_when_cache_is_full(): |
| 103 | manager = EncoderCacheManager(cache_size=10) |
| 104 | |
| 105 | req1 = MockRequest("req1", ["x"], [6]) |
| 106 | req2 = MockRequest("req2", ["y"], [5]) |
| 107 | |
| 108 | assert manager.can_allocate(req1, 0, int(1e9), 0) |
| 109 | manager.allocate(req1, 0) |
| 110 | manager.free_encoder_input(req1, 0) |
| 111 | |
| 112 | assert manager.can_allocate(req2, 0, int(1e9), 0) |
| 113 | manager.allocate(req2, 0) |
| 114 | |
| 115 | # 'x' should have been evicted. |
| 116 | assert "x" not in manager.cached |
| 117 | assert "x" not in manager.num_embeds_by_hash |
| 118 | assert "req1" not in manager.request_cached_ids |
| 119 | assert "x" in manager.get_freed_mm_hashes() |
| 120 | |
| 121 | |
| 122 | def test_get_cached_input_ids(): |
| 123 | manager = EncoderCacheManager(cache_size=10) |
| 124 | req = MockRequest("reqX", ["m", "n", "o"], [2, 4, 3]) |
| 125 | |
| 126 | assert manager.can_allocate(req, 0, int(1e9), 0) |
| 127 | manager.allocate(req, 0) |
| 128 | |
| 129 | assert manager.can_allocate(req, 2, int(1e9), 0) |
| 130 | manager.allocate(req, 2) |
| 131 | |
| 132 | cached_ids = manager.get_cached_input_ids(req) |
| 133 | assert cached_ids == {0, 2} |
| 134 | |
| 135 | |
| 136 | def test_has_cache_restores_from_freeable(): |
| 137 | manager = EncoderCacheManager(cache_size=10) |
| 138 | req = MockRequest("reqY", ["imgZ"], [4]) |
| 139 | |
| 140 | assert manager.can_allocate(req, 0, int(1e9), 0) |
| 141 | manager.allocate(req, 0) |
| 142 | |
| 143 | manager.free_encoder_input(req, 0) |
| 144 | |
| 145 | # Should restore from freeable. |
| 146 | assert manager.check_and_update_cache(req, 0) |
| 147 | assert len(manager.cached["imgZ"]) == 1 |
| 148 | assert manager.request_cached_ids["reqY"] == {0} |
| 149 | assert "imgZ" not in manager.freeable |
| 150 | assert manager.num_freeable_slots == 6 |
| 151 | |
| 152 | |
| 153 | def test_get_freed_mm_hashes_clears_freed_list(): |
| 154 | manager = EncoderCacheManager(cache_size=10) |
| 155 | req1 = MockRequest("reqA", ["a"], [5]) |
| 156 | req2 = MockRequest("reqB", ["b"], [6]) |
| 157 | |
| 158 | assert manager.can_allocate(req1, 0, int(1e9), 0) |
| 159 | manager.allocate(req1, 0) |
| 160 | manager.free_encoder_input(req1, 0) |
| 161 | |
| 162 | # Should trigger eviction of 'a'. |
| 163 | assert manager.can_allocate(req2, 0, int(1e9), 0) |
| 164 | manager.allocate(req2, 0) |
| 165 | |
| 166 | freed = manager.get_freed_mm_hashes() |
| 167 | assert "a" in freed |
| 168 | assert manager.get_freed_mm_hashes() == [] |
| 169 | |
| 170 | |
| 171 | def test_referencing_a_freeable_entry_protects_it_from_eviction(): |
| 172 | """A waiting request can hold an entry it is not scheduled for yet. |
| 173 | |
| 174 | Entries with no referent are evictable. A request that is deferred (e.g. |
| 175 | by an EC connector waiting on another item) still needs the items it |
| 176 | already has, and a connector that hands out one transfer per request |
| 177 | cannot always produce the same item a second time. |
| 178 | """ |
| 179 | manager = EncoderCacheManager(cache_size=10) |
| 180 | owner = MockRequest("owner", ["a"], [5]) |
| 181 | waiter = MockRequest("waiter", ["a"], [5]) |
| 182 | newcomer = MockRequest("newcomer", ["b"], [6]) |
| 183 | |
| 184 | manager.allocate(owner, 0) |
| 185 | manager.free_encoder_input(owner, 0) |
| 186 | assert "a" in manager.freeable |
| 187 | |
| 188 | # The waiting request references it; it is no longer reclaimable. |
| 189 | assert manager.check_and_update_cache(waiter, 0) |
| 190 | assert "a" not in manager.freeable |
| 191 | |
| 192 | # 'b' no longer fits, and 'a' must not be taken away to make room. |
| 193 | assert not manager.can_allocate(newcomer, 0, int(1e9), 0) |
| 194 | assert "a" in manager.cached |
| 195 | assert manager.get_freed_mm_hashes() == [] |
| 196 | |
| 197 | # Once the waiter is done with it, the entry is reclaimable again. |
| 198 | manager.free_encoder_input(waiter, 0) |
| 199 | assert manager.can_allocate(newcomer, 0, int(1e9), 0) |
| 200 | assert manager.get_freed_mm_hashes() == ["a"] |
| 201 | |
| 202 | |
| 203 | def test_reallocated_hash_is_not_reported_as_freed(): |
| 204 | manager = EncoderCacheManager(cache_size=8) |
| 205 | req_a = MockRequest("reqA", ["a"], [4]) |
| 206 | req_b = MockRequest("reqB", ["b"], [4]) |
| 207 | req_c = MockRequest("reqC", ["c"], [4]) |
| 208 | |
| 209 | manager.allocate(req_a, 0) |
| 210 | manager.allocate(req_b, 0) |
| 211 | manager.free(req_a) |
| 212 | manager.free(req_b) |
| 213 | |
| 214 | assert manager.can_allocate(req_c, 0, int(1e9), 0) |
| 215 | manager.allocate(req_c, 0) |
| 216 | assert manager.can_allocate(req_a, 0, int(1e9), 0) |
| 217 | manager.allocate(req_a, 0) |
| 218 | |
| 219 | assert manager.get_freed_mm_hashes() == ["b"] |
| 220 | |
| 221 | |
| 222 | def test_schedule_request_multi_images_respect_space_limit(): |
| 223 | manager = EncoderCacheManager(cache_size=10) |
| 224 | req = MockRequest("reqA", ["a", "b"], [5, 6]) |
| 225 | compute_budget = 100 |
| 226 | |
| 227 | num_tokens_to_schedule = 0 |
| 228 | assert manager.can_allocate(req, 0, compute_budget, num_tokens_to_schedule) |
| 229 | num_tokens_to_schedule += req.get_num_encoder_embeds(0) |
| 230 | compute_budget -= req.get_num_encoder_embeds(0) |
| 231 | |
| 232 | assert not manager.can_allocate(req, 1, compute_budget, num_tokens_to_schedule) |
| 233 | |
| 234 | |
| 235 | def test_schedule_request_multi_images_respect_compute_limit(): |
| 236 | manager = EncoderCacheManager(cache_size=100) |
| 237 | req = MockRequest("reqA", ["a", "b"], [5, 6]) |
| 238 | compute_budget = 10 |
| 239 | num_tokens_to_schedule = 0 |
| 240 | assert manager.can_allocate(req, 0, compute_budget, num_tokens_to_schedule) |
| 241 | num_tokens_to_schedule += req.get_num_encoder_embeds(0) |
| 242 | compute_budget -= req.get_num_encoder_embeds(0) |
| 243 | |
| 244 | assert not manager.can_allocate(req, 1, compute_budget, num_tokens_to_schedule) |
| 245 | |
| 246 | |
| 247 | def test_encoder_cache_with_is_embed_mask(): |
| 248 | class MockRequestWithMask(MockRequest): |
| 249 | def get_num_encoder_embeds(self, input_id: int) -> int: |
| 250 | return self.mm_features[input_id].mm_position.get_num_embeds() |
| 251 | |
| 252 | is_embed = torch.zeros(100, dtype=torch.bool) |
| 253 | is_embed[torch.tensor([5, 15, 25, 35, 45, 55, 65, 75])] = True |
| 254 | |
| 255 | request = MockRequestWithMask("r1", ["img1"], [100]) |
| 256 | request.mm_features[0] = MultiModalFeatureSpec( |
| 257 | data=None, |
| 258 | modality="image", |
| 259 | identifier="img1", |
| 260 | mm_position=PlaceholderRange(offset=0, length=100, is_embed=is_embed), |
| 261 | ) |
| 262 | |
| 263 | manager = EncoderCacheManager(cache_size=100) |
| 264 | manager.allocate(request, 0) |
| 265 | |
| 266 | assert manager.num_free_slots == 92 |
| 267 | assert "img1" in manager.cached |
| 268 | assert manager.request_cached_ids["r1"] == {0} |
| 269 | |
| 270 | old_size = 100 |
| 271 | new_size = request.mm_features[0].mm_position.get_num_embeds() |
| 272 | assert new_size == 8 |
| 273 | savings_ratio = old_size / new_size |
| 274 | assert savings_ratio == 12.5 |
| 275 | |
| 276 | |
| 277 | def test_encoder_cache_tracks_embedding_count_by_identifier(): |
| 278 | class MockRequestWithMask(MockRequest): |
| 279 | def get_num_encoder_embeds(self, input_id: int) -> int: |
| 280 | return self.mm_features[input_id].mm_position.get_num_embeds() |
| 281 | |
| 282 | is_embed = torch.tensor([False, True, True, False, True, False]) |
| 283 | request = MockRequestWithMask("owner", ["img1"], [len(is_embed)]) |
| 284 | request.mm_features[0] = MultiModalFeatureSpec( |
| 285 | data=None, |
| 286 | modality="image", |
| 287 | identifier="img1", |
| 288 | mm_position=PlaceholderRange(offset=0, length=6, is_embed=is_embed), |
| 289 | ) |
| 290 | |
| 291 | manager = EncoderCacheManager(cache_size=10) |
| 292 | manager.allocate(request, 0) |
| 293 | assert manager.get_cached_num_encoder_embeds(request, 0) == 3 |
| 294 | |
| 295 | manager.free(request) |
| 296 | assert manager.get_cached_num_encoder_embeds(request, 0) == 3 |
| 297 | |
| 298 | manager.reset() |
| 299 | assert manager.get_cached_num_encoder_embeds(request, 0) is None |
| 300 | |
| 301 | |
| 302 | def test_encoder_cache_mask_based_retrieval(): |
| 303 | class MockRequestWithMask(MockRequest): |
| 304 | def get_num_encoder_embeds(self, input_id: int) -> int: |
| 305 | return self.mm_features[input_id].mm_position.get_num_embeds() |
| 306 | |
| 307 | is_embed = torch.tensor( |
| 308 | [False, False, True, True, False, True, True, True, False, False] |
| 309 | ) |
| 310 | |
| 311 | request = MockRequestWithMask("r1", ["img1"], [10]) |
| 312 | request.mm_features[0] = MultiModalFeatureSpec( |
| 313 | data=None, |
| 314 | modality="image", |
| 315 | identifier="img1", |
| 316 | mm_position=PlaceholderRange(offset=0, length=10, is_embed=is_embed), |
| 317 | ) |
| 318 | |
| 319 | manager = EncoderCacheManager(cache_size=50) |
| 320 | manager.allocate(request, 0) |
| 321 | |
| 322 | assert request.mm_features[0].mm_position.get_num_embeds() == 5 |
| 323 | |
| 324 | start_idx = 2 |
| 325 | end_idx = 8 |
| 326 | num_embeds_before = is_embed[:start_idx].sum().item() |
| 327 | num_embeds_in_range = is_embed[start_idx:end_idx].sum().item() |
| 328 | |
| 329 | assert num_embeds_before == 0 |
| 330 | assert num_embeds_in_range == 5 |
| 331 | |
| 332 | start_idx = 0 |
| 333 | end_idx = 5 |
| 334 | num_embeds_before = is_embed[:start_idx].sum().item() if start_idx > 0 else 0 |
| 335 | num_embeds_in_range = is_embed[start_idx:end_idx].sum().item() |
| 336 | |
| 337 | assert num_embeds_before == 0 |
| 338 | assert num_embeds_in_range == 2 |
| 339 | |
| 340 | |
| 341 | def test_reset_clears_all_state(): |
| 342 | """Test that reset() clears all cached entries and restores capacity.""" |
| 343 | manager = EncoderCacheManager(cache_size=20) |
| 344 | |
| 345 | req1 = MockRequest("req1", ["img1", "img2"], [5, 3]) |
| 346 | req2 = MockRequest("req2", ["img3"], [4]) |
| 347 | |
| 348 | manager.allocate(req1, 0) |
| 349 | manager.allocate(req1, 1) |
| 350 | manager.allocate(req2, 0) |
| 351 | manager.free_encoder_input(req1, 0) |
| 352 | |
| 353 | req3 = MockRequest("req3", ["img4"], [10]) |
| 354 | manager.free_encoder_input(req1, 1) |
| 355 | manager.free_encoder_input(req2, 0) |
| 356 | manager.can_allocate(req3, 0, int(1e9), 0) |
| 357 | manager.allocate(req3, 0) |
| 358 | |
| 359 | assert len(manager.cached) > 0 |
| 360 | assert manager.num_free_slots < 20 |
| 361 | |
| 362 | manager.reset() |
| 363 | |
| 364 | assert len(manager.cached) == 0 |
| 365 | assert len(manager.num_embeds_by_hash) == 0 |
| 366 | assert len(manager.request_cached_ids) == 0 |
| 367 | assert len(manager.freeable) == 0 |
| 368 | assert len(manager.freed) == 0 |
| 369 | assert manager.num_free_slots == 20 |
| 370 | assert manager.num_freeable_slots == 20 |
| 371 | |
| 372 | |
| 373 | def test_reset_allows_fresh_allocations(): |
| 374 | manager = EncoderCacheManager(cache_size=10) |
| 375 | |
| 376 | req1 = MockRequest("req1", ["img1"], [10]) |
| 377 | manager.allocate(req1, 0) |
| 378 | assert manager.num_free_slots == 0 |
| 379 | |
| 380 | manager.reset() |
| 381 | |
| 382 | req2 = MockRequest("req2", ["img2"], [8]) |
| 383 | assert manager.can_allocate(req2, 0, int(1e9), 0) |
| 384 | manager.allocate(req2, 0) |
| 385 | |
| 386 | assert manager.num_free_slots == 2 |
| 387 | assert "img2" in manager.cached |
| 388 | assert "img1" not in manager.cached |
| 389 | assert manager.request_cached_ids["req2"] == {0} |
| 390 | assert "req1" not in manager.request_cached_ids |
| 391 | |
| 392 | |
| 393 | def test_free_request_with_duplicate_mm_hashes(): |
| 394 | """Freeing a request whose two inputs share the same mm_hash must fully |
| 395 | clean up request_cached_ids. After the first free_encoder_input call, |
| 396 | cached[mm_hash] becomes empty; the second call must still remove the |
| 397 | remaining input_id from request_cached_ids.""" |
| 398 | manager = EncoderCacheManager(cache_size=20) |
| 399 | |
| 400 | req = MockRequest("r1", ["imgA", "imgA"], [4, 4]) |
| 401 | |
| 402 | manager.allocate(req, 0) |
| 403 | # input 1 has the same hash, so it's already cached. |
| 404 | assert manager.check_and_update_cache(req, 1) |
| 405 | assert manager.request_cached_ids["r1"] == {0, 1} |
| 406 | |
| 407 | manager.free(req) |
| 408 | assert "r1" not in manager.request_cached_ids |
| 409 | |
| 410 | |
| 411 | def test_duplicate_mm_hash_stays_referenced_until_last_free(): |
| 412 | """`cached` holds one reference per request, not per position, so freeing |
| 413 | the first of two occurrences must not release the entry.""" |
| 414 | manager = EncoderCacheManager(cache_size=20) |
| 415 | req = MockRequest("r1", ["imgA", "imgA"], [4, 4]) |
| 416 | |
| 417 | manager.allocate(req, 0) |
| 418 | assert manager.check_and_update_cache(req, 1) |
| 419 | |
| 420 | manager.free_encoder_input(req, 0) |
| 421 | assert manager.cached["imgA"] == {"r1"} |
| 422 | assert "imgA" not in manager.freeable |
| 423 | assert manager.num_freeable_slots == 16 |
| 424 | |
| 425 | manager.free_encoder_input(req, 1) |
| 426 | assert not manager.cached["imgA"] |
| 427 | assert manager.freeable["imgA"] == 4 |
| 428 | assert manager.num_freeable_slots == 20 |
| 429 | |
| 430 | |
| 431 | def test_duplicate_mm_hash_is_not_evicted_before_last_use(): |
| 432 | """An item repeated within one request (an image carried across |
| 433 | conversation turns) must not become reclaimable capacity while a later |
| 434 | occurrence still has to be spliced in, or the encoder recomputes it.""" |
| 435 | manager = EncoderCacheManager(cache_size=8) |
| 436 | req = MockRequest("r1", ["imgA", "imgA"], [4, 4]) |
| 437 | other = MockRequest("r2", ["imgB"], [8]) |
| 438 | |
| 439 | manager.allocate(req, 0) |
| 440 | assert manager.check_and_update_cache(req, 1) |
| 441 | manager.free_encoder_input(req, 0) |
| 442 | |
| 443 | assert not manager.can_allocate(other, 0, int(1e9), 0) |
| 444 | assert manager.check_and_update_cache(req, 1) |
| 445 | assert manager.get_freed_mm_hashes() == [] |
| 446 | |
| 447 | |
| 448 | def test_encoder_decoder_cache_manager_reset(): |
| 449 | manager = EncoderDecoderCacheManager(cache_size=20) |
| 450 | |
| 451 | req1 = MockRequest("req1", ["img1"], [5]) |
| 452 | req2 = MockRequest("req2", ["img2"], [3]) |
| 453 | |
| 454 | manager.allocate(req1, 0) |
| 455 | manager.allocate(req2, 0) |
| 456 | manager.free(req1) |
| 457 | manager.get_freed_mm_hashes() |
| 458 | |
| 459 | assert manager.num_free_slots < 20 |
| 460 | |
| 461 | manager.reset() |
| 462 | |
| 463 | assert len(manager.allocated) == 0 |
| 464 | assert len(manager.to_free) == 0 |
| 465 | assert manager.num_free_slots == 20 |
| 466 | |
| 467 | |
| 468 | def test_encoder_decoder_cache_manager_reset_allows_fresh_allocations(): |
| 469 | manager = EncoderDecoderCacheManager(cache_size=10) |
| 470 | |
| 471 | req1 = MockRequest("req1", ["img1"], [10]) |
| 472 | manager.allocate(req1, 0) |
| 473 | assert manager.num_free_slots == 0 |
| 474 | |
| 475 | manager.reset() |
| 476 | |
| 477 | req2 = MockRequest("req2", ["img2"], [8]) |
| 478 | assert manager.can_allocate(req2, 0, int(1e9), 0) |
| 479 | manager.allocate(req2, 0) |
| 480 | |
| 481 | assert manager.num_free_slots == 2 |
| 482 | assert "img2" in manager.allocated |