| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for FileMapper.""" |
| 4 | |
| 5 | from unittest.mock import MagicMock |
| 6 | |
| 7 | from vllm.v1.kv_offload.base import OffloadingSpec, make_offload_key |
| 8 | from vllm.v1.kv_offload.config import ( |
| 9 | OffloadingCacheConfig, |
| 10 | OffloadingConfig, |
| 11 | OffloadingGroupConfig, |
| 12 | OffloadingModelConfig, |
| 13 | OffloadingParallelConfig, |
| 14 | ) |
| 15 | from vllm.v1.kv_offload.file_mapper import FileMapper |
| 16 | |
| 17 | # --------------------------------------------------------------------------- |
| 18 | # Helper |
| 19 | # --------------------------------------------------------------------------- |
| 20 | |
| 21 | |
| 22 | def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: |
| 23 | """Build a FileMapper from a mocked spec carrying a hand-built config.""" |
| 24 | config = OffloadingConfig( |
| 25 | groups=tuple( |
| 26 | OffloadingGroupConfig( |
| 27 | tokens_per_block=tokens_per_block, |
| 28 | layer_names=(layer_name,), |
| 29 | group_id=group_id, |
| 30 | ) |
| 31 | for group_id, (tokens_per_block, layer_name) in enumerate( |
| 32 | kwargs.get("groups", ()) |
| 33 | ) |
| 34 | ), |
| 35 | worker_kv_bytes_per_block=0, |
| 36 | enable_kv_cache_events=False, |
| 37 | extra_config={}, |
| 38 | engine_id="test-engine", |
| 39 | model=OffloadingModelConfig( |
| 40 | name=kwargs.get("model_name", "test-model"), |
| 41 | dtype=kwargs.get("dtype", "float16"), |
| 42 | ), |
| 43 | cache=OffloadingCacheConfig( |
| 44 | tokens_per_hash=kwargs.get("tokens_per_hash", 16), |
| 45 | blocks_per_chunk=kwargs.get("blocks_per_chunk", 1), |
| 46 | ), |
| 47 | parallel=OffloadingParallelConfig( |
| 48 | rank=kwargs.get("rank", 0), |
| 49 | world_size=kwargs.get("world_size", 1), |
| 50 | tp_size=kwargs.get("tp_size", 1), |
| 51 | pp_size=kwargs.get("pp_size", 1), |
| 52 | pcp_size=kwargs.get("pcp_size", 1), |
| 53 | dcp_size=kwargs.get("dcp_size", 1), |
| 54 | data_parallel_index=0, |
| 55 | data_parallel_size=1, |
| 56 | data_parallel_rank_local=None, |
| 57 | is_parallelism_agnostic=kwargs.get("is_parallelism_agnostic", False), |
| 58 | ), |
| 59 | replicated_layout=kwargs.get("replicated_layout", False), |
| 60 | canonical_layout=kwargs.get("canonical_layout", False), |
| 61 | kv_cache_layout=kwargs.get("kv_cache_layout", "LBNHC"), |
| 62 | ) |
| 63 | spec = MagicMock(spec=OffloadingSpec) |
| 64 | spec.config = config |
| 65 | return FileMapper.from_offloading_spec( |
| 66 | root_dir=kwargs.get("root_dir", "/tmp/cache"), |
| 67 | offloading_spec=spec, |
| 68 | blocks_per_file=config.cache.blocks_per_chunk, |
| 69 | parallel_agnostic=kwargs.get("parallel_agnostic", False), |
| 70 | ) |
| 71 | |
| 72 | |
| 73 | # --------------------------------------------------------------------------- |
| 74 | # Tests |
| 75 | # --------------------------------------------------------------------------- |
| 76 | |
| 77 | |
| 78 | def test_get_file_name_full_structure(): |
| 79 | """Path must match: <base_path>_r<rank>/<hhh>/<hh>_g<group_idx>/<hash_hex>.bin. |
| 80 | |
| 81 | Concretely: |
| 82 | - The segment immediately after base_path must end with `_r3` |
| 83 | - The next segment is the first 3 hex chars of the block hash |
| 84 | - The next segment is <2 hex chars>_g<group_idx> |
| 85 | - The final segment is <full hash hex>.bin |
| 86 | """ |
| 87 | rank = 3 |
| 88 | group_idx = 2 |
| 89 | block_hash = bytes(range(8)) # deterministic, non-zero bytes |
| 90 | fm = make_mapper_from_offloading_spec(rank=rank) |
| 91 | key = make_offload_key(block_hash, group_idx) |
| 92 | path = fm.get_file_name(key) |
| 93 | |
| 94 | expected_path = ( |
| 95 | "/tmp/cache/test-model_de3bba26cf36_r3/000/10_g2/0001020304050607.bin" |
| 96 | ) |
| 97 | assert path == expected_path |
| 98 | |
| 99 | |
| 100 | def test_get_run_config_fields(): |
| 101 | fm = make_mapper_from_offloading_spec( |
| 102 | model_name="my-model", |
| 103 | dtype="bfloat16", |
| 104 | tp_size=4, |
| 105 | pp_size=3, |
| 106 | pcp_size=2, |
| 107 | dcp_size=2, |
| 108 | groups=((64, "layer0"),), |
| 109 | tokens_per_hash=64, |
| 110 | blocks_per_chunk=3, |
| 111 | ) |
| 112 | cfg = fm.get_run_config() |
| 113 | assert cfg == { |
| 114 | "model_name": "my-model", |
| 115 | "tokens_per_hash": 64, |
| 116 | "blocks_per_file": 3, |
| 117 | "tp_size": 4, |
| 118 | "pp_size": 3, |
| 119 | "pcp_size": 2, |
| 120 | "dcp_size": 2, |
| 121 | "dtype": "bfloat16", |
| 122 | "kv_cache_groups": [ |
| 123 | { |
| 124 | "tokens_per_block": 64, |
| 125 | "layer_names": ["layer0"], |
| 126 | } |
| 127 | ], |
| 128 | "inference_engine": "vllm", |
| 129 | "parallel_agnostic": False, |
| 130 | } |
| 131 | |
| 132 | |
| 133 | def test_get_config_file_path(): |
| 134 | fm = make_mapper_from_offloading_spec() |
| 135 | config_path = fm.get_config_file_path() |
| 136 | assert config_path == f"{fm.base_path}/config.json" |
| 137 | |
| 138 | |
| 139 | def test_hybrid_file_identity_uses_resolved_tokens_per_hash(): |
| 140 | # For heterogeneous groups the namespace records the resolved hash |
| 141 | # granularity (GCD of the group block sizes), which is the actual |
| 142 | # granularity of the offload block hashes. |
| 143 | fm = make_mapper_from_offloading_spec( |
| 144 | groups=((12, "full_layer"), (16, "mla_layer")), |
| 145 | tokens_per_hash=4, |
| 146 | ) |
| 147 | assert fm.fields["tokens_per_hash"] == 4 |
| 148 | assert fm.fields["kv_cache_groups"] == [ |
| 149 | {"tokens_per_block": 12, "layer_names": ["full_layer"]}, |
| 150 | {"tokens_per_block": 16, "layer_names": ["mla_layer"]}, |
| 151 | ] |
| 152 | |
| 153 | |
| 154 | # --------------------------------------------------------------------------- |
| 155 | # parallel_agnostic: opt-in honored only when the config marks the layout |
| 156 | # parallelism-agnostic (predicate computation is covered in test_factory.py) |
| 157 | # --------------------------------------------------------------------------- |
| 158 | |
| 159 | |
| 160 | def test_parallel_agnostic_collapses_namespace_when_config_allows(): |
| 161 | fm = make_mapper_from_offloading_spec( |
| 162 | tp_size=4, |
| 163 | pp_size=3, |
| 164 | pcp_size=2, |
| 165 | dcp_size=2, |
| 166 | rank=1, |
| 167 | is_parallelism_agnostic=True, |
| 168 | parallel_agnostic=True, |
| 169 | ) |
| 170 | assert fm.fields["tp_size"] == 1 |
| 171 | assert fm.fields["pp_size"] == 1 |
| 172 | assert fm.fields["pcp_size"] == 1 |
| 173 | assert fm.fields["dcp_size"] == 1 |
| 174 | assert fm.rank == 0 |
| 175 | assert "parallel_agnostic" not in fm.fields |
| 176 | |
| 177 | |
| 178 | def test_parallel_agnostic_ignored_when_config_disallows(): |
| 179 | fm = make_mapper_from_offloading_spec( |
| 180 | tp_size=2, |
| 181 | rank=1, |
| 182 | is_parallelism_agnostic=False, |
| 183 | parallel_agnostic=True, |
| 184 | ) |
| 185 | assert fm.fields["tp_size"] == 2 |
| 186 | assert fm.rank == 1 |
| 187 | assert fm.fields["parallel_agnostic"] is False |
| 188 | |
| 189 | |
| 190 | def test_namespace_kept_without_parallel_agnostic_opt_in(): |
| 191 | fm = make_mapper_from_offloading_spec( |
| 192 | tp_size=2, |
| 193 | rank=1, |
| 194 | is_parallelism_agnostic=True, |
| 195 | parallel_agnostic=False, |
| 196 | ) |
| 197 | assert fm.fields["tp_size"] == 2 |
| 198 | assert fm.rank == 1 |
| 199 | assert fm.fields["parallel_agnostic"] is False |
| 200 | |
| 201 | |
| 202 | def test_parallel_agnostic_separates_persistent_layouts(): |
| 203 | agnostic = make_mapper_from_offloading_spec( |
| 204 | is_parallelism_agnostic=True, |
| 205 | parallel_agnostic=True, |
| 206 | ) |
| 207 | specific = make_mapper_from_offloading_spec( |
| 208 | is_parallelism_agnostic=False, |
| 209 | parallel_agnostic=True, |
| 210 | ) |
| 211 | |
| 212 | assert agnostic.base_path != specific.base_path |
| 213 | assert "parallel_agnostic" not in agnostic.fields |
| 214 | assert specific.fields["parallel_agnostic"] is False |
| 215 | |
| 216 | |
| 217 | def test_canonical_layout_changes_storage_namespace(): |
| 218 | # Canonical bytes are not interchangeable with the direct layout, so the |
| 219 | # format id must fork the storage namespace. |
| 220 | direct = make_mapper_from_offloading_spec() |
| 221 | canonical = make_mapper_from_offloading_spec(canonical_layout=True) |
| 222 | assert "canonical_format" not in direct.fields |
| 223 | assert canonical.fields["canonical_format"] == "v1-nhd" |
| 224 | assert direct.base_path != canonical.base_path |
| 225 | |
| 226 | |
| 227 | # --------------------------------------------------------------------------- |
| 228 | # replicated_layout: OR'd into parallel-agnostic identity for compact rows |
| 229 | # --------------------------------------------------------------------------- |
| 230 | |
| 231 | |
| 232 | def test_replicated_layout_collapses_parallel_identity(): |
| 233 | shared = dict( |
| 234 | model_name="mla-model", |
| 235 | groups=((16, "mla_layer"),), |
| 236 | replicated_layout=True, |
| 237 | parallel_agnostic=True, |
| 238 | ) |
| 239 | tp2 = make_mapper_from_offloading_spec(tp_size=2, world_size=2, rank=1, **shared) |
| 240 | tp4 = make_mapper_from_offloading_spec(tp_size=4, world_size=4, rank=3, **shared) |
| 241 | |
| 242 | assert tp2.base_path == tp4.base_path |
| 243 | for fm in (tp2, tp4): |
| 244 | assert fm.fields["tp_size"] == 1 |
| 245 | assert fm.fields["pp_size"] == 1 |
| 246 | assert fm.fields["pcp_size"] == 1 |
| 247 | assert fm.fields["dcp_size"] == 1 |
| 248 | assert fm.rank == 0 |
| 249 | assert "parallel_agnostic" not in fm.fields |
| 250 | assert fm.fields["replicated_layout"] is True |
| 251 | assert fm.get_file_name(make_offload_key(b"\x01" * 8, 0)).startswith( |
| 252 | f"{fm.base_path}_r0/" |
| 253 | ) |
| 254 | |
| 255 | |
| 256 | def test_replicated_layout_requires_caller_opt_in(): |
| 257 | fm = make_mapper_from_offloading_spec( |
| 258 | tp_size=2, |
| 259 | world_size=2, |
| 260 | rank=1, |
| 261 | replicated_layout=True, |
| 262 | parallel_agnostic=False, |
| 263 | ) |
| 264 | assert fm.fields["tp_size"] == 2 |
| 265 | assert fm.rank == 1 |
| 266 | assert fm.fields["parallel_agnostic"] is False |
| 267 | assert "replicated_layout" not in fm.fields |
| 268 | baseline = make_mapper_from_offloading_spec( |
| 269 | tp_size=2, |
| 270 | world_size=2, |
| 271 | rank=1, |
| 272 | replicated_layout=False, |
| 273 | parallel_agnostic=False, |
| 274 | ) |
| 275 | assert fm.base_path == baseline.base_path |
| 276 | |
| 277 | |
| 278 | def test_non_replicated_keeps_parallel_identity(): |
| 279 | fm = make_mapper_from_offloading_spec( |
| 280 | tp_size=4, |
| 281 | world_size=4, |
| 282 | rank=2, |
| 283 | replicated_layout=False, |
| 284 | is_parallelism_agnostic=False, |
| 285 | parallel_agnostic=True, |
| 286 | ) |
| 287 | assert fm.fields["tp_size"] == 4 |
| 288 | assert fm.rank == 2 |
| 289 | assert fm.fields["parallel_agnostic"] is False |
| 290 | assert fm.get_file_name(make_offload_key(b"\x02" * 8, 0)).startswith( |
| 291 | f"{fm.base_path}_r2/" |
| 292 | ) |
| 293 | |
| 294 | |
| 295 | def test_replicated_and_parallelism_agnostic_separate_layouts(): |
| 296 | shared = dict( |
| 297 | model_name="shared-model", |
| 298 | groups=((16, "layer0"),), |
| 299 | tp_size=2, |
| 300 | world_size=2, |
| 301 | rank=1, |
| 302 | parallel_agnostic=True, |
| 303 | ) |
| 304 | via_agnostic = make_mapper_from_offloading_spec( |
| 305 | is_parallelism_agnostic=True, |
| 306 | replicated_layout=False, |
| 307 | **shared, |
| 308 | ) |
| 309 | via_replicated = make_mapper_from_offloading_spec( |
| 310 | is_parallelism_agnostic=False, |
| 311 | replicated_layout=True, |
| 312 | **shared, |
| 313 | ) |
| 314 | assert via_agnostic.base_path != via_replicated.base_path |
| 315 | assert "replicated_layout" not in via_agnostic.fields |
| 316 | assert via_replicated.fields["replicated_layout"] is True |
| 317 | |
| 318 | |
| 319 | def test_replicated_layout_run_config_tp_invariant(): |
| 320 | shared = dict( |
| 321 | model_name="mla-model", |
| 322 | groups=((16, "mla_layer"),), |
| 323 | replicated_layout=True, |
| 324 | parallel_agnostic=True, |
| 325 | ) |
| 326 | tp2 = make_mapper_from_offloading_spec(tp_size=2, world_size=2, rank=0, **shared) |
| 327 | tp4 = make_mapper_from_offloading_spec(tp_size=4, world_size=4, rank=2, **shared) |
| 328 | assert tp2.get_run_config() == tp4.get_run_config() |