| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for the HIP/ROCm paths in vllm.v1.simple_kv_offload.cuda_mem_ops. |
| 4 | |
| 5 | Deliberately not skipped off ROCm: the version gate takes a plain version int |
| 6 | and the library loader is monkeypatched, so nothing here needs a HIP runtime |
| 7 | and the logic stays covered by the default (NVIDIA) CI that runs this |
| 8 | directory. |
| 9 | """ |
| 10 | |
| 11 | import ctypes |
| 12 | |
| 13 | import numpy as np |
| 14 | import pytest |
| 15 | |
| 16 | from vllm.v1.simple_kv_offload import cuda_mem_ops |
| 17 | from vllm.v1.simple_kv_offload.cuda_mem_ops import ( |
| 18 | BatchMemcpyParams, |
| 19 | _CUmemcpyAttributes, |
| 20 | _load_hip_runtime, |
| 21 | _num_attrs_for_hip_version, |
| 22 | copy_blocks, |
| 23 | ) |
| 24 | |
| 25 | |
| 26 | # hipMemcpyBatchAsync rejects numAttrs > 0 on ROCm 7.2.1-7.2.3 and accepts it |
| 27 | # on 7.13+. HIP encodes the version as major*10_000_000 + minor*100_000 + patch. |
| 28 | @pytest.mark.parametrize( |
| 29 | ("version", "expected"), |
| 30 | [ |
| 31 | (0, 0), # unknown (hipRuntimeGetVersion failed) -> conservative 0 |
| 32 | (70_100_000, 0), # ROCm 7.1 |
| 33 | (70_200_000, 0), # ROCm 7.2.0 |
| 34 | (70_253_211, 0), # ROCm 7.2.2 (patch 53211) |
| 35 | (70_300_000, 0), # ROCm 7.3 |
| 36 | (71_200_000, 0), # ROCm 7.12 -- just below the 7.13 cutoff |
| 37 | (71_300_000, 1), # ROCm 7.13 -- first release accepting numAttrs > 0 |
| 38 | (71_326_193, 1), # ROCm 7.13 nightly (patch 26193) |
| 39 | (71_400_000, 1), # ROCm 7.14 |
| 40 | (71_500_000, 1), # ROCm 7.15 |
| 41 | (80_000_000, 1), # ROCm 8.0 |
| 42 | ], |
| 43 | ) |
| 44 | def test_num_attrs_for_hip_version(version: int, expected: int): |
| 45 | assert _num_attrs_for_hip_version(version) == expected |
| 46 | |
| 47 | |
| 48 | @pytest.mark.parametrize( |
| 49 | ("available", "expected"), |
| 50 | [ |
| 51 | # Devel install: the unversioned symlink is present and preferred. |
| 52 | ({"libamdhip64.so", "libamdhip64.so.7"}, "libamdhip64.so"), |
| 53 | # Runtime-only / wheel-packaged ROCm: versioned soname only. |
| 54 | ({"libamdhip64.so.7"}, "libamdhip64.so.7"), |
| 55 | ({"libamdhip64.so.6"}, "libamdhip64.so.6"), |
| 56 | ], |
| 57 | ) |
| 58 | def test_load_hip_runtime_falls_back_to_versioned_soname( |
| 59 | monkeypatch, available: set[str], expected: str |
| 60 | ): |
| 61 | loaded = [] |
| 62 | |
| 63 | def fake_cdll(name, mode=0): |
| 64 | if name not in available: |
| 65 | raise OSError(f"{name}: cannot open shared object file") |
| 66 | loaded.append(name) |
| 67 | return name |
| 68 | |
| 69 | monkeypatch.setattr(cuda_mem_ops.ctypes, "CDLL", fake_cdll) |
| 70 | assert _load_hip_runtime() == expected |
| 71 | assert loaded == [expected] |
| 72 | |
| 73 | |
| 74 | def test_load_hip_runtime_reports_every_candidate_when_all_fail(monkeypatch): |
| 75 | def fake_cdll(name, mode=0): |
| 76 | raise OSError("cannot open shared object file") |
| 77 | |
| 78 | monkeypatch.setattr(cuda_mem_ops.ctypes, "CDLL", fake_cdll) |
| 79 | with pytest.raises(OSError, match="could not load the HIP runtime") as exc: |
| 80 | _load_hip_runtime() |
| 81 | for name in ("libamdhip64.so", "libamdhip64.so.7", "libamdhip64.so.6"): |
| 82 | assert name in str(exc.value) |
| 83 | |
| 84 | |
| 85 | def _fake_params(num_layers: int, bytes_per_block: int = 512) -> BatchMemcpyParams: |
| 86 | """Minimal params for driving copy_blocks() without a real device. |
| 87 | |
| 88 | Bases are 0 and the batch-memcpy function is monkeypatched, so no memory is |
| 89 | ever dereferenced; only the descriptor bookkeeping (count/chunking) matters. |
| 90 | """ |
| 91 | return BatchMemcpyParams( |
| 92 | src_bases=np.zeros(num_layers, dtype=np.uint64), |
| 93 | dst_bases=np.zeros(num_layers, dtype=np.uint64), |
| 94 | bpb=np.full(num_layers, bytes_per_block, dtype=np.uint64), |
| 95 | num_layers=num_layers, |
| 96 | attrs=_CUmemcpyAttributes(), |
| 97 | attrs_idx=ctypes.c_size_t(0), |
| 98 | num_attrs=0, |
| 99 | fail_idx=ctypes.c_size_t(0), |
| 100 | stream_handle=0, |
| 101 | ) |
| 102 | |
| 103 | |
| 104 | def _record_counts(monkeypatch) -> list[int]: |
| 105 | """Patch the batch-memcpy call to record the descriptor count per call.""" |
| 106 | counts: list[int] = [] |
| 107 | |
| 108 | def fake_fn(dst, src, sizes, count, *rest): |
| 109 | counts.append(count) |
| 110 | return 0 |
| 111 | |
| 112 | monkeypatch.setattr(cuda_mem_ops, "_batch_memcpy", (fake_fn, 0)) |
| 113 | return counts |
| 114 | |
| 115 | |
| 116 | # ROCm chunks the descriptor batch at the cap (hipMemcpyBatchAsync faults above |
| 117 | # 8192/call); CUDA leaves the cap at 0 and issues a single call. |
| 118 | @pytest.mark.parametrize( |
| 119 | ("max_desc", "num_layers", "num_blocks", "expected_counts"), |
| 120 | [ |
| 121 | (8192, 3, 5000, [8192, 6808]), # 15000 desc -> 8192 + 6808 |
| 122 | (8192, 1, 8192, [8192]), # exactly at the cap -> single call |
| 123 | (8192, 1, 8193, [8192, 1]), # one over -> spills into a second call |
| 124 | (4096, 61, 135, [4096, 4096, 43]), # bs=1 Kimi shape (8235 desc) |
| 125 | (0, 61, 135, [8235]), # CUDA uncapped -> single call |
| 126 | ], |
| 127 | ) |
| 128 | def test_copy_blocks_chunks_at_descriptor_cap( |
| 129 | monkeypatch, max_desc, num_layers, num_blocks, expected_counts |
| 130 | ): |
| 131 | counts = _record_counts(monkeypatch) |
| 132 | monkeypatch.setattr(cuda_mem_ops, "_max_batch_descriptors", max_desc) |
| 133 | |
| 134 | ids = list(range(num_blocks)) |
| 135 | copy_blocks(ids, ids, _fake_params(num_layers)) |
| 136 | |
| 137 | assert counts == expected_counts |
| 138 | assert sum(counts) == num_layers * num_blocks |
| 139 | |
| 140 | |
| 141 | def test_copy_blocks_noop_on_empty(monkeypatch): |
| 142 | counts = _record_counts(monkeypatch) |
| 143 | monkeypatch.setattr(cuda_mem_ops, "_max_batch_descriptors", 8192) |
| 144 | copy_blocks([], [], _fake_params(num_layers=4)) |
| 145 | assert counts == [] |