| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """test that invalid blocks are evicted from prefix cache to prevent pollution. |
| 5 | |
| 6 | verifies that when sync-loading fails, invalid blocks are removed from the |
| 7 | prefix cache hash table so future requests cannot match and reuse corrupted data. |
| 8 | """ |
| 9 | |
| 10 | from collections.abc import Callable |
| 11 | from unittest.mock import Mock |
| 12 | |
| 13 | import pytest |
| 14 | |
| 15 | from vllm.v1.core.sched.scheduler import Scheduler |
| 16 | from vllm.v1.request import Request, RequestStatus |
| 17 | |
| 18 | from .utils import ( |
| 19 | create_model_runner_output, |
| 20 | create_request, |
| 21 | create_scheduler, |
| 22 | create_vllm_config, |
| 23 | ) |
| 24 | |
| 25 | pytestmark = pytest.mark.cpu_test |
| 26 | |
| 27 | |
| 28 | def _make_get_num_new_matched_tokens( |
| 29 | req_num_new_matched_tokens: dict[str, int], |
| 30 | async_load: bool, |
| 31 | ) -> Callable[[Request, int], tuple[int, bool]]: |
| 32 | def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]: |
| 33 | value = req_num_new_matched_tokens.get(request.request_id, 0) |
| 34 | return value, async_load |
| 35 | |
| 36 | return get_num_new_matched_tokens |
| 37 | |
| 38 | |
| 39 | @pytest.fixture |
| 40 | def fail_scheduler(): |
| 41 | """Scheduler with kv_load_failure_policy='fail'.""" |
| 42 | vllm_config = create_vllm_config() |
| 43 | vllm_config.kv_transfer_config.kv_load_failure_policy = "fail" |
| 44 | return create_scheduler(vllm_config) |
| 45 | |
| 46 | |
| 47 | def test_invalid_blocks_evicted_prevents_cache_pollution( |
| 48 | fail_scheduler: Scheduler, |
| 49 | ): |
| 50 | """Verify invalid blocks are evicted to prevent future cache hits. |
| 51 | |
| 52 | scenario: |
| 53 | 1. request 1 loads externally-computed blocks (sync mode) |
| 54 | 2. some blocks fail to load and are marked invalid |
| 55 | 3. with fail policy, invalid blocks should be evicted from prefix cache |
| 56 | 4. request is marked as FINISHED_ERROR |
| 57 | """ |
| 58 | num_prompt_blocks = 100 |
| 59 | num_external_computed_blocks = 99 |
| 60 | invalid_block_idx = 50 |
| 61 | |
| 62 | num_prompt_tokens = num_prompt_blocks * fail_scheduler.block_size |
| 63 | num_external_computed_tokens = ( |
| 64 | num_external_computed_blocks * fail_scheduler.block_size |
| 65 | ) |
| 66 | |
| 67 | # request 1: will have invalid blocks |
| 68 | request1 = create_request(num_tokens=num_prompt_tokens, request_id=1) |
| 69 | fail_scheduler.add_request(request=request1) |
| 70 | |
| 71 | req_num_new_matched_tokens = { |
| 72 | request1.request_id: num_external_computed_tokens, |
| 73 | } |
| 74 | |
| 75 | # mock connector indicating sync load |
| 76 | fail_scheduler.connector = Mock() |
| 77 | fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( |
| 78 | _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False) |
| 79 | ) |
| 80 | fail_scheduler.connector.request_finished.return_value = (False, None) |
| 81 | fail_scheduler.connector.take_events.return_value = () |
| 82 | |
| 83 | scheduler_output = fail_scheduler.schedule() |
| 84 | |
| 85 | # request should be running with sync KV load |
| 86 | assert len(fail_scheduler.running) == 1 |
| 87 | assert request1.status == RequestStatus.RUNNING |
| 88 | |
| 89 | # get allocated block IDs |
| 90 | req_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0] |
| 91 | invalid_block_id = req_block_ids[invalid_block_idx] |
| 92 | invalid_block_ids = {invalid_block_id} |
| 93 | |
| 94 | # get the block object to verify eviction later |
| 95 | block = fail_scheduler.kv_cache_manager.block_pool.blocks[invalid_block_id] |
| 96 | |
| 97 | # cache the blocks to simulate they've been computed and cached |
| 98 | # (in real scenario blocks would be cached after compute) |
| 99 | fail_scheduler.kv_cache_manager.cache_blocks(request1, num_external_computed_tokens) |
| 100 | |
| 101 | # verify block has a hash (is cached) before reporting invalid blocks |
| 102 | assert block.block_hash is not None, ( |
| 103 | f"block {invalid_block_id} should be cached (have a hash) before " |
| 104 | f"eviction test, but hash is None" |
| 105 | ) |
| 106 | |
| 107 | # report invalid blocks |
| 108 | model_runner_output = create_model_runner_output( |
| 109 | [request1], |
| 110 | invalid_block_ids=invalid_block_ids, |
| 111 | use_eos=False, |
| 112 | ) |
| 113 | |
| 114 | fail_scheduler.update_from_output(scheduler_output, model_runner_output) |
| 115 | |
| 116 | # verify request finished with error (fail policy) |
| 117 | assert request1.status == RequestStatus.FINISHED_ERROR |
| 118 | |
| 119 | # critical assertion: invalid block and all subsequent blocks should be evicted |
| 120 | # all blocks from invalid_block_idx onwards become invalid since they were |
| 121 | # computed based on the failed block |
| 122 | for idx in range(invalid_block_idx, len(req_block_ids)): |
| 123 | block_id = req_block_ids[idx] |
| 124 | block_obj = fail_scheduler.kv_cache_manager.block_pool.blocks[block_id] |
| 125 | assert block_obj.block_hash is None, ( |
| 126 | f"block {block_id} at index {idx} should have been evicted " |
| 127 | f"(hash reset to None), but hash is {block_obj.block_hash}. " |
| 128 | f"All blocks from index {invalid_block_idx} onwards should be evicted " |
| 129 | f"since they depend on the invalid block at index {invalid_block_idx}." |
| 130 | ) |
| 131 | |
| 132 | # verify cache contains exactly the valid blocks (before first affected block) |
| 133 | # and none of the invalid blocks (from first affected block onwards) |
| 134 | |
| 135 | # valid blocks: all blocks before invalid_block_idx should be cached |
| 136 | for idx in range(invalid_block_idx): |
| 137 | block_id = req_block_ids[idx] |
| 138 | block_obj = fail_scheduler.kv_cache_manager.block_pool.blocks[block_id] |
| 139 | assert block_obj.block_hash is not None, ( |
| 140 | f"valid block {block_id} at index {idx} should still be cached " |
| 141 | f"(have a hash), but hash is None. Only blocks from index " |
| 142 | f"{invalid_block_idx} onwards should be evicted." |
| 143 | ) |
| 144 | |
| 145 | # invalid blocks: verify they're not in the cached_block_hash_to_block map |
| 146 | cached_blocks = ( |
| 147 | fail_scheduler.kv_cache_manager.block_pool.cached_block_hash_to_block |
| 148 | ) |
| 149 | cached_block_ids = { |
| 150 | b.block_id |
| 151 | for blocks_val in cached_blocks._cache.values() |
| 152 | for b in ( |
| 153 | [blocks_val] if not isinstance(blocks_val, dict) else blocks_val.values() |
| 154 | ) |
| 155 | } |
| 156 | |
| 157 | for idx in range(invalid_block_idx, len(req_block_ids)): |
| 158 | block_id = req_block_ids[idx] |
| 159 | assert block_id not in cached_block_ids, ( |
| 160 | f"invalid block {block_id} at index {idx} should not be in cache hash table" |
| 161 | ) |