| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit test for the Triton ``swap_blocks_batch`` fast-path kernel.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.platforms import current_platform |
| 9 | from vllm.v1.kv_offload.cpu.swap_blocks_triton import swap_blocks_batch |
| 10 | |
| 11 | |
| 12 | def _addrs(buffers: list[torch.Tensor]) -> torch.Tensor: |
| 13 | return torch.tensor([b.data_ptr() for b in buffers], dtype=torch.int64) |
| 14 | |
| 15 | |
| 16 | @pytest.mark.skipif( |
| 17 | not current_platform.is_cuda(), reason="Triton swap fast path requires CUDA" |
| 18 | ) |
| 19 | def test_triton_swap_copies_source_bytes(): |
| 20 | # 8-byte-aligned, sub-threshold sizes covering 8 KiB chunk boundaries and |
| 21 | # odd tail-mask lengths, with enough descriptors to take the Triton path. |
| 22 | sizes = [8, 4096, 8192, 8200, 16384, 4088] * 8 |
| 23 | src = [torch.randint(256, (s,), dtype=torch.uint8, device="cuda") for s in sizes] |
| 24 | dst = [torch.zeros_like(s) for s in src] |
| 25 | sizes_t = torch.tensor(sizes, dtype=torch.int64) |
| 26 | |
| 27 | swap_blocks_batch(_addrs(src), _addrs(dst), sizes_t.clone(), bytes_per_chunk=8192) |
| 28 | torch.accelerator.synchronize() |
| 29 | |
| 30 | for s, t in zip(src, dst): |
| 31 | assert torch.equal(t, s) # kernel copied the source bytes verbatim |