| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import time |
| 4 | import uuid |
| 5 | |
| 6 | import numpy as np |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.v1.kv_offload.base import ( |
| 11 | CanonicalKVCacheRef, |
| 12 | CanonicalPageMapping, |
| 13 | CopyRun, |
| 14 | GPULoadStoreSpec, |
| 15 | ) |
| 16 | from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec |
| 17 | from vllm.v1.kv_offload.cpu.gpu_worker import ( |
| 18 | SingleDirectionOffloadingHandler, |
| 19 | _build_copy_plan, |
| 20 | _canonical_block_sizes, |
| 21 | _canonical_page_ids, |
| 22 | pin_mmap_region, |
| 23 | ) |
| 24 | from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion |
| 25 | |
| 26 | |
| 27 | def _ref(mapping: CanonicalPageMapping, tensor_idx: int = 0) -> CanonicalKVCacheRef: |
| 28 | return CanonicalKVCacheRef( |
| 29 | tensor_idx=tensor_idx, |
| 30 | page_size_bytes=mapping.local_page_size_bytes, |
| 31 | mapping=mapping, |
| 32 | ) |
| 33 | |
| 34 | |
| 35 | def _nhd_mapping() -> CanonicalPageMapping: |
| 36 | # 4-token NHD page at tp=4, rank 2: K and V runs of 4x128B fragments |
| 37 | runs = ( |
| 38 | CopyRun(0, 256, 128, 4, 128, 512), |
| 39 | CopyRun(512, 2304, 128, 4, 128, 512), |
| 40 | ) |
| 41 | return CanonicalPageMapping(4096, 1024, runs, 1, 0, True) |
| 42 | |
| 43 | |
| 44 | def test_copy_plan_unrolls_runs(): |
| 45 | plan = _build_copy_plan(_ref(_nhd_mapping()), gpu_to_cpu=True) |
| 46 | k_dst = [256, 768, 1280, 1792] |
| 47 | assert plan.frag_offsets_src.tolist() == [0, 128, 256, 384, 512, 640, 768, 896] |
| 48 | assert plan.frag_offsets_dst.tolist() == k_dst + [2048 + o for o in k_dst] |
| 49 | assert plan.frag_sizes.tolist() == [128] * 8 |
| 50 | assert plan.num_frags == 8 |
| 51 | |
| 52 | |
| 53 | def test_load_direction_swaps_offsets(): |
| 54 | store = _build_copy_plan(_ref(_nhd_mapping()), gpu_to_cpu=True) |
| 55 | load = _build_copy_plan(_ref(_nhd_mapping()), gpu_to_cpu=False) |
| 56 | assert np.array_equal(store.frag_offsets_src, load.frag_offsets_dst) |
| 57 | assert np.array_equal(store.frag_offsets_dst, load.frag_offsets_src) |
| 58 | assert np.array_equal(store.frag_sizes, load.frag_sizes) |
| 59 | |
| 60 | |
| 61 | def test_writer_rotation_matches_is_writer(): |
| 62 | # Replicas take turns writing shared canonical pages, keyed by the |
| 63 | # CPU-side canonical page id; the enumeration must agree with is_writer |
| 64 | identity = CopyRun(0, 0, 256, 1, 256, 256) |
| 65 | mapping = CanonicalPageMapping(256, 256, (identity,), 2, 1, True) |
| 66 | ids = _canonical_page_ids( |
| 67 | np.array([3, 7, 9]), blocks_per_chunk=4, count=10, skip_count=2 |
| 68 | ) |
| 69 | assert ids.tolist() == [14, 15, 28, 29, 30, 31, 36, 37, 38, 39] |
| 70 | mask = ids % mapping.num_writers == mapping.writer_index |
| 71 | assert mask.tolist() == [mapping.is_writer(int(b)) for b in ids] |
| 72 | |
| 73 | |
| 74 | def test_canonical_block_sizes_take_max_per_tensor(): |
| 75 | identity = CopyRun(0, 0, 512, 1, 512, 512) |
| 76 | small = CanonicalPageMapping(2048, 512, (identity,), 1, 0, False) |
| 77 | refs = [[_ref(_nhd_mapping(), 0), _ref(small, 0)], [_ref(small, 1)]] |
| 78 | assert _canonical_block_sizes(refs, 2) == [4096, 2048] |
| 79 | |
| 80 | |
| 81 | def _tp2_rank_mapping(rank: int) -> CanonicalPageMapping: |
| 82 | # 4-token NHD page, 4 total heads of 64B, tp=2: rank holds 2 heads, |
| 83 | # so K and V each scatter as 4 per-token 128B fragments |
| 84 | runs = ( |
| 85 | CopyRun(0, rank * 128, 128, 4, 128, 256), |
| 86 | CopyRun(512, 1024 + rank * 128, 128, 4, 128, 256), |
| 87 | ) |
| 88 | return CanonicalPageMapping(2048, 1024, runs, 1, 0, True) |
| 89 | |
| 90 | |
| 91 | def _whole_page_mapping() -> CanonicalPageMapping: |
| 92 | identity = CopyRun(0, 0, 2048, 1, 2048, 2048) |
| 93 | return CanonicalPageMapping(2048, 2048, (identity,), 1, 0, True) |
| 94 | |
| 95 | |
| 96 | def _transfer(handler, num_blocks: int, gpu_to_cpu: bool) -> None: |
| 97 | block_ids = list(range(num_blocks)) |
| 98 | gpu_spec = GPULoadStoreSpec( |
| 99 | block_ids, group_sizes=(num_blocks,), block_indices=(0,) |
| 100 | ) |
| 101 | cpu_spec = CPULoadStoreSpec(block_ids) |
| 102 | src, dst = (gpu_spec, cpu_spec) if gpu_to_cpu else (cpu_spec, gpu_spec) |
| 103 | assert handler.transfer_async(0, src, dst) |
| 104 | deadline = time.time() + 30 |
| 105 | while time.time() < deadline: |
| 106 | if handler.get_finished(): |
| 107 | return |
| 108 | time.sleep(0.001) |
| 109 | raise TimeoutError("transfer did not complete") |
| 110 | |
| 111 | |
| 112 | def _canonical_handler(gpu_tensor, cpu_tensor, mapping, gpu_to_cpu): |
| 113 | page = mapping.local_page_size_bytes |
| 114 | refs = [[CanonicalKVCacheRef(tensor_idx=0, page_size_bytes=page, mapping=mapping)]] |
| 115 | return SingleDirectionOffloadingHandler( |
| 116 | gpu_tensors=[gpu_tensor], |
| 117 | cpu_tensors=[cpu_tensor], |
| 118 | blocks_per_chunk=1, |
| 119 | layer_refs_per_group=refs, |
| 120 | gpu_to_cpu=gpu_to_cpu, |
| 121 | canonical_layout=True, |
| 122 | ) |
| 123 | |
| 124 | |
| 125 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 126 | def test_gpu_roundtrip_assembles_canonical_page_across_ranks(): |
| 127 | """Two TP2-style rank handlers must scatter into one shared canonical |
| 128 | CPU page such that each rank reloads bit-exact and a whole-page (TP1) |
| 129 | reader sees both shards — the cross-topology contract.""" |
| 130 | torch.manual_seed(0) |
| 131 | num_blocks = 4 |
| 132 | gpu_rank = [ |
| 133 | torch.randint(-128, 128, (num_blocks, 1024), dtype=torch.int8, device="cuda") |
| 134 | for _ in range(2) |
| 135 | ] |
| 136 | cpu_canonical = torch.zeros(num_blocks, 2048, dtype=torch.int8, pin_memory=True) |
| 137 | |
| 138 | for rank in (0, 1): |
| 139 | store = _canonical_handler( |
| 140 | gpu_rank[rank], cpu_canonical, _tp2_rank_mapping(rank), gpu_to_cpu=True |
| 141 | ) |
| 142 | _transfer(store, num_blocks, gpu_to_cpu=True) |
| 143 | torch.accelerator.synchronize() |
| 144 | |
| 145 | # independent oracle: replay each rank's runs in numpy |
| 146 | expected = np.zeros((num_blocks, 2048), dtype=np.int8) |
| 147 | for rank in (0, 1): |
| 148 | local = gpu_rank[rank].cpu().numpy() |
| 149 | for run in _tp2_rank_mapping(rank).runs: |
| 150 | for i in range(run.num_fragments): |
| 151 | lo = run.local_offset + i * run.local_stride |
| 152 | co = run.canonical_offset + i * run.canonical_stride |
| 153 | expected[:, co : co + run.fragment_size] = local[ |
| 154 | :, lo : lo + run.fragment_size |
| 155 | ] |
| 156 | assert np.array_equal(cpu_canonical.numpy(), expected) |
| 157 | |
| 158 | # each rank reloads its shard bit-exact |
| 159 | gpu_back = torch.zeros(num_blocks, 1024, dtype=torch.int8, device="cuda") |
| 160 | load = _canonical_handler( |
| 161 | gpu_back, cpu_canonical, _tp2_rank_mapping(0), gpu_to_cpu=False |
| 162 | ) |
| 163 | _transfer(load, num_blocks, gpu_to_cpu=False) |
| 164 | torch.accelerator.synchronize() |
| 165 | assert torch.equal(gpu_back, gpu_rank[0]) |
| 166 | |
| 167 | # a whole-page reader (TP1 topology) sees the assembled page |
| 168 | gpu_full = torch.zeros(num_blocks, 2048, dtype=torch.int8, device="cuda") |
| 169 | load_full = _canonical_handler( |
| 170 | gpu_full, cpu_canonical, _whole_page_mapping(), gpu_to_cpu=False |
| 171 | ) |
| 172 | _transfer(load_full, num_blocks, gpu_to_cpu=False) |
| 173 | torch.accelerator.synchronize() |
| 174 | assert torch.equal(gpu_full.cpu(), cpu_canonical) |
| 175 | |
| 176 | |
| 177 | # 4-token NHD page, 4 total heads of 64 bytes: canonical page holds |
| 178 | # [K: token x head][V: token x head] = 2048 bytes per block |
| 179 | _TOTAL_HEADS = 4 |
| 180 | _HEAD_BYTES = 64 |
| 181 | _BLOCK_TOKENS = 4 |
| 182 | _CANONICAL_PAGE = 2 * _BLOCK_TOKENS * _TOTAL_HEADS * _HEAD_BYTES |
| 183 | |
| 184 | |
| 185 | def _nhd_shard_mapping(tp: int, rank: int) -> CanonicalPageMapping: |
| 186 | """Rank's head-shard mapping into the canonical NHD page at the given tp.""" |
| 187 | local_heads = _TOTAL_HEADS // tp |
| 188 | frag = local_heads * _HEAD_BYTES |
| 189 | canonical_row = _TOTAL_HEADS * _HEAD_BYTES |
| 190 | k_run = CopyRun(0, rank * frag, frag, _BLOCK_TOKENS, frag, canonical_row) |
| 191 | v_run = CopyRun( |
| 192 | _BLOCK_TOKENS * frag, |
| 193 | _BLOCK_TOKENS * canonical_row + rank * frag, |
| 194 | frag, |
| 195 | _BLOCK_TOKENS, |
| 196 | frag, |
| 197 | canonical_row, |
| 198 | ) |
| 199 | return CanonicalPageMapping( |
| 200 | _CANONICAL_PAGE, 2 * _BLOCK_TOKENS * frag, (k_run, v_run), 1, 0, True |
| 201 | ) |
| 202 | |
| 203 | |
| 204 | def _head_shard(full_kv: torch.Tensor, tp: int, rank: int) -> torch.Tensor: |
| 205 | """This rank's local page rows out of the (blocks, 2, tokens, heads, |
| 206 | head_bytes) ground truth.""" |
| 207 | local_heads = _TOTAL_HEADS // tp |
| 208 | shard = full_kv[:, :, :, rank * local_heads : (rank + 1) * local_heads, :] |
| 209 | return shard.reshape(full_kv.shape[0], -1).contiguous() |
| 210 | |
| 211 | |
| 212 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 213 | @pytest.mark.parametrize("writer_tp,reader_tp", [(2, 4), (4, 2), (2, 1), (4, 4)]) |
| 214 | def test_cross_topology_roundtrip(writer_tp: int, reader_tp: int): |
| 215 | """KV written at one tp must be readable at another: writer ranks scatter |
| 216 | head shards into a shared canonical region, reader ranks gather their own |
| 217 | shards, and every reader must see the writers' ground-truth bytes.""" |
| 218 | torch.manual_seed(0) |
| 219 | num_blocks = 3 |
| 220 | full_kv = torch.randint( |
| 221 | -128, |
| 222 | 128, |
| 223 | (num_blocks, 2, _BLOCK_TOKENS, _TOTAL_HEADS, _HEAD_BYTES), |
| 224 | dtype=torch.int8, |
| 225 | ) |
| 226 | |
| 227 | engine_id = str(uuid.uuid4()) |
| 228 | row_stride = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT |
| 229 | assert row_stride >= _CANONICAL_PAGE |
| 230 | regions: list[SharedOffloadRegion] = [] |
| 231 | |
| 232 | def canonical_view(rank: int, world_size: int) -> torch.Tensor: |
| 233 | region = SharedOffloadRegion( |
| 234 | engine_id=engine_id, |
| 235 | num_chunks=num_blocks, |
| 236 | rank=rank, |
| 237 | kv_bytes_per_chunk=row_stride, |
| 238 | cpu_page_size=row_stride // world_size, |
| 239 | ) |
| 240 | regions.append(region) |
| 241 | # The Triton load path dereferences CPU pointers on the GPU, which is |
| 242 | # only legal on pinned memory; production pins via CPUOffloadingWorker |
| 243 | pin_mmap_region(region) |
| 244 | assert region.is_pinned |
| 245 | return region.create_next_canonical_view(_CANONICAL_PAGE) |
| 246 | |
| 247 | try: |
| 248 | for rank in range(writer_tp): |
| 249 | store = _canonical_handler( |
| 250 | _head_shard(full_kv, writer_tp, rank).cuda(), |
| 251 | canonical_view(rank, writer_tp), |
| 252 | _nhd_shard_mapping(writer_tp, rank), |
| 253 | gpu_to_cpu=True, |
| 254 | ) |
| 255 | _transfer(store, num_blocks, gpu_to_cpu=True) |
| 256 | torch.accelerator.synchronize() |
| 257 | |
| 258 | for rank in range(reader_tp): |
| 259 | expected = _head_shard(full_kv, reader_tp, rank) |
| 260 | gpu_out = torch.zeros_like(expected, device="cuda") |
| 261 | load = _canonical_handler( |
| 262 | gpu_out, |
| 263 | canonical_view(rank, reader_tp), |
| 264 | _nhd_shard_mapping(reader_tp, rank), |
| 265 | gpu_to_cpu=False, |
| 266 | ) |
| 267 | _transfer(load, num_blocks, gpu_to_cpu=False) |
| 268 | torch.accelerator.synchronize() |
| 269 | assert torch.equal(gpu_out.cpu(), expected), ( |
| 270 | f"reader tp={reader_tp} rank={rank} bytes diverge from the " |
| 271 | f"tp={writer_tp} writers' ground truth" |
| 272 | ) |
| 273 | finally: |
| 274 | for region in regions: |
| 275 | region.cleanup() |