| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | |
| 6 | from vllm.model_executor.offloader.prefetch import _get_next_prefetch_index |
| 7 | |
| 8 | pytestmark = pytest.mark.skip_global_cleanup |
| 9 | |
| 10 | |
| 11 | @pytest.mark.parametrize( |
| 12 | ("module_count", "prefetch_step", "expected_targets"), |
| 13 | [ |
| 14 | (3, 1, [1, 2, 0]), |
| 15 | (3, 2, [2, 1, 0]), |
| 16 | (4, 2, [2, 3, 0, 1]), |
| 17 | (5, 2, [2, 3, 4, 1, 0]), |
| 18 | (3, 3, [0, 1, 2]), |
| 19 | (3, 4, [1, 2, 0]), |
| 20 | ], |
| 21 | ) |
| 22 | def test_next_prefetch_index_preserves_slot_ownership( |
| 23 | module_count: int, |
| 24 | prefetch_step: int, |
| 25 | expected_targets: list[int], |
| 26 | ) -> None: |
| 27 | targets = [ |
| 28 | _get_next_prefetch_index(index, prefetch_step, module_count) |
| 29 | for index in range(module_count) |
| 30 | ] |
| 31 | |
| 32 | assert targets == expected_targets |
| 33 | if prefetch_step < module_count: |
| 34 | assert all( |
| 35 | target % prefetch_step == index % prefetch_step |
| 36 | for index, target in enumerate(targets) |
| 37 | ) |