| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for contiguous KV cache packing. |
| 4 | |
| 5 | Every cache group packs its layers densely into one block; groups overlay each other |
| 6 | (a block ID is owned by one group at a time), so the packed block stride is the |
| 7 | largest group's packing. The layout decides whether the layer dim sits outside the |
| 8 | block dim (a contiguous region per layer) or inside it (all layers' pages within each |
| 9 | block); the allocation is the same either way. |
| 10 | """ |
| 11 | |
| 12 | from dataclasses import replace |
| 13 | from math import lcm |
| 14 | from unittest.mock import MagicMock |
| 15 | |
| 16 | import pytest |
| 17 | import torch |
| 18 | |
| 19 | from vllm.config import CacheConfig |
| 20 | from vllm.v1.core.kv_cache_manager import KVCacheManager |
| 21 | from vllm.v1.core.kv_cache_utils import ( |
| 22 | _get_kv_cache_bytes_per_block, |
| 23 | _get_packed_kv_cache_groups, |
| 24 | _pool_bytes_per_block, |
| 25 | generate_scheduler_kv_cache_config, |
| 26 | get_kv_cache_config_from_groups, |
| 27 | get_kv_cache_groups, |
| 28 | resolve_kv_cache_block_sizes, |
| 29 | ) |
| 30 | from vllm.v1.hisparse.layout import _build_hisparse_kv_cache_tensors |
| 31 | from vllm.v1.kv_cache_interface import ( |
| 32 | CircularBufferSpec, |
| 33 | FullAttentionSpec, |
| 34 | HiSparseHotSpec, |
| 35 | KVCacheGroupSpec, |
| 36 | KVCacheLayout, |
| 37 | KVCacheSpec, |
| 38 | MambaSpec, |
| 39 | MLAAttentionSpec, |
| 40 | SlidingWindowMLASpec, |
| 41 | UniformTypeKVCacheSpecs, |
| 42 | iter_layer_specs, |
| 43 | ) |
| 44 | from vllm.v1.worker.utils import allocate_kv_cache |
| 45 | |
| 46 | MEMORY = 8 * 1024 * 1024 |
| 47 | |
| 48 | |
| 49 | def _mla(head_size: int) -> MLAAttentionSpec: |
| 50 | return MLAAttentionSpec( |
| 51 | block_size=64, num_kv_heads=1, head_size=head_size, dtype=torch.uint8 |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | def _full() -> FullAttentionSpec: |
| 56 | return FullAttentionSpec( |
| 57 | block_size=16, num_kv_heads=2, head_size=64, dtype=torch.float16 |
| 58 | ) |
| 59 | |
| 60 | |
| 61 | def _uniform_group(specs: dict) -> KVCacheGroupSpec: |
| 62 | return KVCacheGroupSpec( |
| 63 | list(specs), |
| 64 | UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=specs), |
| 65 | ) |
| 66 | |
| 67 | |
| 68 | def _mixed_page_groups(n_mla=3, n_idx=3, n_swa=5): |
| 69 | """A DeepSeek-V4-style hybrid: two groups with different page mixes.""" |
| 70 | g1 = {f"mla.{i}": _mla(512) for i in range(n_mla)} |
| 71 | g1.update({f"idx.{i}": _mla(128) for i in range(n_idx)}) |
| 72 | g2 = {f"swa.{i}": _mla(512) for i in range(n_swa)} |
| 73 | return [_uniform_group(g1), _uniform_group(g2)], g1, g2 |
| 74 | |
| 75 | |
| 76 | def _mock_vllm_config(layout: str | None): |
| 77 | config = MagicMock() |
| 78 | config.cache_config = CacheConfig() |
| 79 | config.cache_config.num_gpu_blocks_override = None |
| 80 | config.cache_config.kv_cache_layout = layout |
| 81 | config.attention_config.hisparse_config = None |
| 82 | return config |
| 83 | |
| 84 | |
| 85 | def test_packed_alignment_preserves_hot_pages_and_generic_stride(): |
| 86 | """One shared stride must satisfy both constraints, regardless of rounding order.""" |
| 87 | spec = replace(_mla(128), block_stride_alignment=1152) |
| 88 | hot = HiSparseHotSpec(block_size=16, page_size=5120, blocks_per_request=2) |
| 89 | groups = [KVCacheGroupSpec(["attention"], spec), KVCacheGroupSpec(["hot"], hot)] |
| 90 | alignment = lcm(1152, hot.page_size_bytes) |
| 91 | stride = _pool_bytes_per_block(groups) |
| 92 | assert stride == alignment |
| 93 | layout = KVCacheLayout.BLHNC |
| 94 | tensors = _build_hisparse_kv_cache_tensors(groups, 3, 3 * stride, layout, stride) |
| 95 | assert {tensor.block_stride for tensor in tensors} == {alignment} |
| 96 | assert {tensor.size for tensor in tensors} == {3 * alignment} |
| 97 | |
| 98 | |
| 99 | @pytest.mark.parametrize("spec", [_full(), _mla(128)]) |
| 100 | def test_block_alignment_survives_merge_and_rejects_incompatible_specs(spec): |
| 101 | aligned = replace(spec, block_stride_alignment=512) |
| 102 | assert type(spec).merge([aligned, aligned]).block_stride_alignment == 512 |
| 103 | with pytest.raises(AssertionError): |
| 104 | type(spec).merge([aligned, spec]) |
| 105 | with pytest.raises(ValueError, match="must be positive"): |
| 106 | replace(spec, block_stride_alignment=0) |
| 107 | |
| 108 | |
| 109 | def _pages(groups) -> dict[str, int]: |
| 110 | return { |
| 111 | name: group.kv_cache_spec.kv_cache_specs[name].page_size_bytes |
| 112 | if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) |
| 113 | else group.kv_cache_spec.page_size_bytes |
| 114 | for group in groups |
| 115 | for name in group.layer_names |
| 116 | } |
| 117 | |
| 118 | |
| 119 | def _expected_bytes_per_block(groups) -> int: |
| 120 | pages = _pages(groups) |
| 121 | return max(sum(pages[n] for n in g.layer_names) for g in groups) |
| 122 | |
| 123 | |
| 124 | def _bind(config, layout: str): |
| 125 | return allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout[layout], None) |
| 126 | |
| 127 | |
| 128 | MAIN_KV_PAGE_BYTES = 2_048 |
| 129 | COMPRESSED_PAGE_BYTES = 128 |
| 130 | NUM_CACHE_TUPLES = 3 |
| 131 | BYTES_PER_BLOCK = NUM_CACHE_TUPLES * (MAIN_KV_PAGE_BYTES + COMPRESSED_PAGE_BYTES) |
| 132 | |
| 133 | |
| 134 | def _main_kv_name(layer_index: int) -> str: |
| 135 | return f"model.layers.{layer_index}.self_attn" |
| 136 | |
| 137 | |
| 138 | def _compressed_name(layer_index: int) -> str: |
| 139 | return f"model.layers.{layer_index}.self_attn.indexer.compressed_key_cache" |
| 140 | |
| 141 | |
| 142 | def _compressor_state_name(layer_index: int) -> str: |
| 143 | return f"model.layers.{layer_index}.self_attn.indexer.raw_key_cache" |
| 144 | |
| 145 | |
| 146 | def _make_csa_linear_specs( |
| 147 | num_mamba: int = 7, |
| 148 | num_tuples: int = NUM_CACHE_TUPLES, |
| 149 | ) -> dict[str, KVCacheSpec]: |
| 150 | specs: dict[str, KVCacheSpec] = {} |
| 151 | for layer_index in range(num_mamba): |
| 152 | specs[f"model.layers.{layer_index}.linear_attn"] = MambaSpec( |
| 153 | block_size=16, |
| 154 | shapes=((32,),), |
| 155 | dtypes=(torch.bfloat16,), |
| 156 | ) |
| 157 | specs["model.layers.0.ple"] = MambaSpec( |
| 158 | block_size=16, |
| 159 | shapes=((24,),), |
| 160 | dtypes=(torch.bfloat16,), |
| 161 | tp_replicated=True, |
| 162 | ) |
| 163 | for layer_index in range(num_tuples): |
| 164 | specs[_main_kv_name(layer_index)] = FullAttentionSpec( |
| 165 | block_size=16, |
| 166 | num_kv_heads=2, |
| 167 | head_size=16, |
| 168 | head_size_v=16, |
| 169 | dtype=torch.bfloat16, |
| 170 | ) |
| 171 | specs[_compressed_name(layer_index)] = MLAAttentionSpec( |
| 172 | block_size=16, |
| 173 | num_kv_heads=1, |
| 174 | head_size=16, |
| 175 | dtype=torch.bfloat16, |
| 176 | tokens_per_state=4, |
| 177 | ) |
| 178 | specs[_compressor_state_name(layer_index)] = CircularBufferSpec( |
| 179 | block_size=4, |
| 180 | num_kv_heads=1, |
| 181 | head_size=8, |
| 182 | head_size_v=0, |
| 183 | dtype=torch.bfloat16, |
| 184 | ) |
| 185 | return specs |
| 186 | |
| 187 | |
| 188 | def _shared_layout_config(): |
| 189 | config = _mock_vllm_config("BLNHC") |
| 190 | config.scheduler_config.disable_hybrid_kv_cache_manager = False |
| 191 | config.speculative_config = None |
| 192 | config.model_config.max_model_len = 16 |
| 193 | config.model_config.get_total_num_hidden_layers.return_value = 64 |
| 194 | config.model_config.get_total_num_kv_heads.return_value = 2 |
| 195 | config.model_config.get_num_kv_heads.return_value = 2 |
| 196 | config.parallel_config.pipeline_parallel_size = 1 |
| 197 | config.parallel_config.decode_context_parallel_size = 1 |
| 198 | config.cache_config.block_size = 16 |
| 199 | config.cache_config.enable_prefix_caching = False |
| 200 | config.cache_config.prefix_match_unit = None |
| 201 | config.cache_config.mamba_cache_mode = "none" |
| 202 | return config |
| 203 | |
| 204 | |
| 205 | class TestCSALinearGrouping: |
| 206 | """A CSA + linear-attention model (sparse attention with a compressor ring, |
| 207 | plus sharded GDN and one TP-replicated PLE state) goes through the generic |
| 208 | packed-group path; no model-specific branch is involved.""" |
| 209 | |
| 210 | @staticmethod |
| 211 | def _mamba_groups(groups): |
| 212 | """(group, per-layer mamba spec) for every group holding mamba state.""" |
| 213 | out = [] |
| 214 | for group in groups: |
| 215 | spec = next(iter(iter_layer_specs(group.kv_cache_spec))) |
| 216 | if isinstance(spec, MambaSpec): |
| 217 | out.append((group, spec)) |
| 218 | return out |
| 219 | |
| 220 | def test_replicated_state_gets_its_own_group(self): |
| 221 | """The PLE state is TP-replicated and a different size, so it must not |
| 222 | share a manager group with the sharded GDN states (the NIXL worker |
| 223 | requires exactly one single-layer replicated group).""" |
| 224 | groups = get_kv_cache_groups(_shared_layout_config(), _make_csa_linear_specs()) |
| 225 | |
| 226 | mamba_groups = self._mamba_groups(groups) |
| 227 | replicated = [g for g, spec in mamba_groups if spec.tp_replicated] |
| 228 | assert len(replicated) == 1 |
| 229 | assert replicated[0].layer_names == ["model.layers.0.ple"] |
| 230 | sharded = [g for g, spec in mamba_groups if not spec.tp_replicated] |
| 231 | assert sharded, "GDN states must keep their own groups" |
| 232 | assert sorted(n for g in sharded for n in g.layer_names) == sorted( |
| 233 | f"model.layers.{i}.linear_attn" for i in range(7) |
| 234 | ) |
| 235 | |
| 236 | def test_roles_land_in_separate_groups(self): |
| 237 | groups = get_kv_cache_groups(_shared_layout_config(), _make_csa_linear_specs()) |
| 238 | |
| 239 | owner = next(g for g in groups if _main_kv_name(0) in g.layer_names) |
| 240 | assert sorted(owner.layer_names) == sorted( |
| 241 | [ |
| 242 | *(_main_kv_name(i) for i in range(NUM_CACHE_TUPLES)), |
| 243 | *(_compressed_name(i) for i in range(NUM_CACHE_TUPLES)), |
| 244 | ] |
| 245 | ) |
| 246 | assert owner.kv_cache_spec.prefix_cacheable |
| 247 | |
| 248 | scratch = next(g for g in groups if _compressor_state_name(0) in g.layer_names) |
| 249 | assert scratch.layer_names == [ |
| 250 | _compressor_state_name(i) for i in range(NUM_CACHE_TUPLES) |
| 251 | ] |
| 252 | assert not scratch.kv_cache_spec.prefix_cacheable |
| 253 | |
| 254 | def test_every_group_fits_one_packed_block(self): |
| 255 | config = _shared_layout_config() |
| 256 | groups = get_kv_cache_groups(config, _make_csa_linear_specs()) |
| 257 | bytes_per_block = _get_kv_cache_bytes_per_block(groups) |
| 258 | |
| 259 | pages = _pages(groups) |
| 260 | for group in groups: |
| 261 | assert sum(pages[n] for n in group.layer_names) <= bytes_per_block |
| 262 | |
| 263 | kv_cache_config = get_kv_cache_config_from_groups( |
| 264 | config, groups, available_memory=bytes_per_block * 32 |
| 265 | ) |
| 266 | assert kv_cache_config.num_blocks == 32 |
| 267 | # Groups overlay from byte 0 of each block, so a layer never addresses |
| 268 | # past the block it belongs to. |
| 269 | for tensor in kv_cache_config.kv_cache_tensors: |
| 270 | assert tensor.block_stride == bytes_per_block |
| 271 | assert tensor.offset < bytes_per_block |
| 272 | |
| 273 | def test_scratch_group_survives_computed_block_truncation(self): |
| 274 | """The scratch group contributes no computed blocks, so truncating a |
| 275 | lookup result must skip it: its block size is the ring capacity, which |
| 276 | neither divides the hit length nor bounds the (empty) block list.""" |
| 277 | config = _shared_layout_config() |
| 278 | config.cache_config.enable_prefix_caching = True |
| 279 | groups = get_kv_cache_groups(config, _make_csa_linear_specs()) |
| 280 | kv_cache_config = get_kv_cache_config_from_groups( |
| 281 | config, groups, available_memory=BYTES_PER_BLOCK * 64 |
| 282 | ) |
| 283 | scheduler_config = generate_scheduler_kv_cache_config([kv_cache_config]) |
| 284 | manager = KVCacheManager( |
| 285 | scheduler_config, |
| 286 | max_model_len=8192, |
| 287 | enable_caching=True, |
| 288 | hash_block_size=16, |
| 289 | scheduler_block_size=16, |
| 290 | ) |
| 291 | blocks = manager.create_kv_cache_blocks( |
| 292 | tuple( |
| 293 | manager.block_pool.get_new_blocks(3) |
| 294 | if group.kv_cache_spec.prefix_cacheable |
| 295 | else [] |
| 296 | for group in scheduler_config.kv_cache_groups |
| 297 | ) |
| 298 | ) |
| 299 | truncated = manager.truncate_computed_blocks(blocks, 48) |
| 300 | scratch_index = next( |
| 301 | i |
| 302 | for i, group in enumerate(scheduler_config.kv_cache_groups) |
| 303 | if not group.kv_cache_spec.prefix_cacheable |
| 304 | ) |
| 305 | assert truncated.blocks[scratch_index] == [] |
| 306 | |
| 307 | def test_prefix_hits_respect_compression_alignment(self): |
| 308 | config = _shared_layout_config() |
| 309 | config.cache_config.enable_prefix_caching = True |
| 310 | config.cache_config.prefix_match_unit = 16 |
| 311 | config.cache_config.mamba_cache_mode = "align" |
| 312 | specs = { |
| 313 | name: replace(spec, mamba_cache_mode="align") |
| 314 | if isinstance(spec, MambaSpec) |
| 315 | else spec |
| 316 | for name, spec in _make_csa_linear_specs(num_tuples=1).items() |
| 317 | } |
| 318 | groups = get_kv_cache_groups(config, specs) |
| 319 | kv_cache_config = get_kv_cache_config_from_groups( |
| 320 | config, groups, available_memory=8 * MAIN_KV_PAGE_BYTES |
| 321 | ) |
| 322 | assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (16, 16) |
| 323 | |
| 324 | # 2 tokens is not a multiple of the compression ratio 4: a prefix hit |
| 325 | # could land inside a partially filled compressed state. |
| 326 | config.cache_config.prefix_match_unit = 2 |
| 327 | with pytest.raises(ValueError, match="per-state compression"): |
| 328 | resolve_kv_cache_block_sizes(kv_cache_config, config) |
| 329 | |
| 330 | def test_scratch_ring_does_not_drag_hash_granularity(self): |
| 331 | config = _shared_layout_config() |
| 332 | config.cache_config.enable_prefix_caching = True |
| 333 | config.cache_config.mamba_cache_mode = "align" |
| 334 | specs = { |
| 335 | name: replace(spec, mamba_cache_mode="align") |
| 336 | if isinstance(spec, MambaSpec) |
| 337 | else spec |
| 338 | for name, spec in _make_csa_linear_specs(num_tuples=1).items() |
| 339 | } |
| 340 | groups = get_kv_cache_groups(config, specs) |
| 341 | kv_cache_config = get_kv_cache_config_from_groups( |
| 342 | config, groups, available_memory=8 * MAIN_KV_PAGE_BYTES |
| 343 | ) |
| 344 | # The hash granularity is the GCD over prefix-cacheable groups only; |
| 345 | # the 4-token scratch ring is excluded (it would drag it to 4). |
| 346 | assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (16, 16) |
| 347 | |
| 348 | def test_compressed_attention_hashes_can_be_finer_than_cache_hits(self): |
| 349 | config = _shared_layout_config() |
| 350 | config.cache_config.enable_prefix_caching = True |
| 351 | specs = { |
| 352 | "compressed.4": MLAAttentionSpec( |
| 353 | block_size=256, |
| 354 | num_kv_heads=1, |
| 355 | head_size=16, |
| 356 | dtype=torch.bfloat16, |
| 357 | tokens_per_state=4, |
| 358 | ), |
| 359 | "compressed.128": MLAAttentionSpec( |
| 360 | block_size=256, |
| 361 | num_kv_heads=1, |
| 362 | head_size=16, |
| 363 | dtype=torch.bfloat16, |
| 364 | tokens_per_state=128, |
| 365 | ), |
| 366 | "compressor_state.4": SlidingWindowMLASpec( |
| 367 | block_size=4, |
| 368 | num_kv_heads=1, |
| 369 | head_size=8, |
| 370 | head_size_v=0, |
| 371 | dtype=torch.bfloat16, |
| 372 | sliding_window=8, |
| 373 | ), |
| 374 | } |
| 375 | groups = get_kv_cache_groups(config, specs) |
| 376 | kv_cache_config = get_kv_cache_config_from_groups( |
| 377 | config, groups, available_memory=8 * MAIN_KV_PAGE_BYTES |
| 378 | ) |
| 379 | |
| 380 | # Hashes are computed every 4 tokens, but without an align-mode Mamba |
| 381 | # group cache hits remain on the 256-token scheduler boundary. |
| 382 | assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (256, 4) |
| 383 | |
| 384 | @pytest.mark.parametrize("wide", ["unbalanced_attention", "unsplittable_state"]) |
| 385 | def test_mamba_split_measures_the_block_the_other_groups_already_force(self, wide): |
| 386 | """Whatever fixes the block stride -- a bucket with unequal layer counts |
| 387 | per page size, which is emitted whole, or a single-layer state that |
| 388 | cannot be split at all -- the mamba layers must be sized against it. |
| 389 | Sizing them against a narrower bucket splits them past what the block |
| 390 | already fits, spending a pool block per extra group for no saving.""" |
| 391 | config = _mock_vllm_config("BLNHC") |
| 392 | config.speculative_config = None |
| 393 | specs = {} |
| 394 | if wide == "unbalanced_attention": |
| 395 | for i in range(3): |
| 396 | specs[f"wide.{i}"] = _mla(1024) |
| 397 | specs["wide.odd"] = _mla(800) |
| 398 | else: |
| 399 | specs["wide.state"] = MambaSpec( |
| 400 | block_size=16, |
| 401 | shapes=((51_200,),), |
| 402 | dtypes=(torch.bfloat16,), |
| 403 | tp_replicated=True, |
| 404 | ) |
| 405 | # Mixed and balanced, but far narrower than the bucket above. |
| 406 | for i in range(10): |
| 407 | specs[f"narrow.a.{i}"] = FullAttentionSpec( |
| 408 | block_size=16, num_kv_heads=1, head_size=8, dtype=torch.uint8 |
| 409 | ) |
| 410 | specs[f"narrow.b.{i}"] = FullAttentionSpec( |
| 411 | block_size=16, num_kv_heads=1, head_size=16, dtype=torch.uint8 |
| 412 | ) |
| 413 | for i in range(100): |
| 414 | specs[f"gdn.{i}"] = MambaSpec( |
| 415 | block_size=16, shapes=((512,),), dtypes=(torch.bfloat16,) |
| 416 | ) |
| 417 | |
| 418 | groups = _get_packed_kv_cache_groups(config, specs) |
| 419 | gdn = [g for g in groups if g.layer_names[0].startswith("gdn.")] |
| 420 | |
| 421 | assert _get_kv_cache_bytes_per_block(groups) == sum( |
| 422 | specs[name].page_size_bytes for name in specs if name.startswith("wide.") |
| 423 | ) |
| 424 | # That block holds every GDN state at once, so the repeat pattern alone |
| 425 | # decides the split; the cap must not add groups on top of it. |
| 426 | assert len(gdn) == 10 |
| 427 | |
| 428 | |
| 429 | class TestSlidingWindowBucketCap: |
| 430 | def test_sliding_window_bucket_is_capped_at_the_main_page(self): |
| 431 | """DeepSeek-V4.1 shape: 43 SlidingWindowMLASpec SWA caches beside an |
| 432 | unbalanced 8-layer paged MLA bucket. Left whole, the SWA bucket would |
| 433 | set the block width and the MLA group would fill under a third of |
| 434 | every block; capped, it splits into groups no wider than the MLA |
| 435 | page.""" |
| 436 | config = _shared_layout_config() |
| 437 | specs: dict[str, KVCacheSpec] = {} |
| 438 | for layer, ratio in ((2, 2), (8, 2), (14, 2), (20, 1)): |
| 439 | specs[f"layers.{layer}.attn"] = MLAAttentionSpec( |
| 440 | block_size=128, |
| 441 | num_kv_heads=1, |
| 442 | head_size=584, |
| 443 | dtype=torch.uint8, |
| 444 | tokens_per_state=ratio, |
| 445 | alignment=576, |
| 446 | ) |
| 447 | specs[f"layers.{layer}.idx"] = MLAAttentionSpec( |
| 448 | block_size=128, |
| 449 | num_kv_heads=1, |
| 450 | head_size=132, |
| 451 | dtype=torch.uint8, |
| 452 | tokens_per_state=ratio, |
| 453 | alignment=576, |
| 454 | ) |
| 455 | for layer in range(43): |
| 456 | specs[f"layers.{layer}.swa"] = SlidingWindowMLASpec( |
| 457 | block_size=64, |
| 458 | num_kv_heads=1, |
| 459 | head_size=584, |
| 460 | dtype=torch.uint8, |
| 461 | sliding_window=128, |
| 462 | alignment=576, |
| 463 | ) |
| 464 | groups = get_kv_cache_groups(config, specs) |
| 465 | pages = _pages(groups) |
| 466 | main = next(g for g in groups if "layers.2.attn" in g.layer_names) |
| 467 | main_bytes = sum(pages[n] for n in main.layer_names) |
| 468 | assert _get_kv_cache_bytes_per_block(groups) == main_bytes |
| 469 | swa_groups = [g for g in groups if g.layer_names[0].endswith(".swa")] |
| 470 | per_group = main_bytes // pages["layers.0.swa"] |
| 471 | assert len(swa_groups) == -(-43 // per_group) |
| 472 | for g in swa_groups: |
| 473 | assert sum(pages[n] for n in g.layer_names) <= main_bytes |
| 474 | assert sorted(n for g in swa_groups for n in g.layer_names) == sorted( |
| 475 | f"layers.{i}.swa" for i in range(43) |
| 476 | ) |
| 477 | |
| 478 | |
| 479 | class TestDensePacking: |
| 480 | def test_bytes_per_block_is_largest_group(self): |
| 481 | groups, g1, g2 = _mixed_page_groups() |
| 482 | assert _get_kv_cache_bytes_per_block(groups) == _expected_bytes_per_block( |
| 483 | groups |
| 484 | ) |
| 485 | |
| 486 | config = get_kv_cache_config_from_groups( |
| 487 | _mock_vllm_config("BLHNC"), groups, MEMORY |
| 488 | ) |
| 489 | # Groups overlay: both start at offset 0. |
| 490 | assert [tensor.offset for tensor in config.kv_cache_tensors].count(0) == 2 |
| 491 | assert [tensor.layers for tensor in config.kv_cache_tensors] == [ |
| 492 | list(g1)[:3], |
| 493 | list(g1)[3:], |
| 494 | list(g2), |
| 495 | ] |
| 496 | |
| 497 | def test_layers_within_a_group_are_dense(self): |
| 498 | groups, _, _ = _mixed_page_groups() |
| 499 | pages = _pages(groups) |
| 500 | config = get_kv_cache_config_from_groups( |
| 501 | _mock_vllm_config("BLHNC"), groups, MEMORY |
| 502 | ) |
| 503 | offsets = { |
| 504 | name: tensor.offset + i * tensor.layer_stride |
| 505 | for tensor in config.kv_cache_tensors |
| 506 | for i, name in enumerate(tensor.layers) |
| 507 | } |
| 508 | for group in groups: |
| 509 | expected = 0 |
| 510 | for name in group.layer_names: |
| 511 | assert offsets[name] == expected |
| 512 | expected += pages[name] |
| 513 | |
| 514 | @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"]) |
| 515 | def test_allocation_is_layout_invariant(self, layout): |
| 516 | specs = {f"l.{i}": _full() for i in range(4)} |
| 517 | groups = [KVCacheGroupSpec(list(specs), _full())] |
| 518 | config = get_kv_cache_config_from_groups( |
| 519 | _mock_vllm_config(layout), groups, MEMORY |
| 520 | ) |
| 521 | (tensor,) = config.kv_cache_tensors |
| 522 | page = _full().page_size_bytes |
| 523 | assert tensor.layers == list(specs) |
| 524 | assert config.num_blocks == MEMORY // (4 * page) |
| 525 | assert tensor.size == 4 * page * config.num_blocks |
| 526 | if layout == "LBNHC": |
| 527 | assert (tensor.layer_stride, tensor.block_stride) == ( |
| 528 | page * config.num_blocks, |
| 529 | page, |
| 530 | ) |
| 531 | else: |
| 532 | assert (tensor.layer_stride, tensor.block_stride) == (page, 4 * page) |
| 533 | |
| 534 | @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"]) |
| 535 | def test_single_group_mixed_pages_follows_layout(self, layout): |
| 536 | specs = {"mla.0": _mla(512), "mla.1": _mla(512), "idx.0": _mla(128)} |
| 537 | groups = [_uniform_group(specs)] |
| 538 | config = get_kv_cache_config_from_groups( |
| 539 | _mock_vllm_config(layout), groups, MEMORY |
| 540 | ) |
| 541 | block_stride = _expected_bytes_per_block(groups) |
| 542 | mla_tensor, idx_tensor = config.kv_cache_tensors |
| 543 | assert mla_tensor.layers == ["mla.0", "mla.1"] |
| 544 | assert idx_tensor.layers == ["idx.0"] |
| 545 | assert {t.size for t in config.kv_cache_tensors} == { |
| 546 | block_stride * config.num_blocks |
| 547 | } |
| 548 | if layout == "LBNHC": |
| 549 | assert ( |
| 550 | idx_tensor.offset == 2 * _mla(512).page_size_bytes * config.num_blocks |
| 551 | ) |
| 552 | assert mla_tensor.block_stride == _mla(512).page_size_bytes |
| 553 | else: |
| 554 | assert idx_tensor.offset == 2 * _mla(512).page_size_bytes |
| 555 | assert mla_tensor.block_stride == block_stride |
| 556 | |
| 557 | def test_overlaid_groups_alias_and_stay_isolated(self): |
| 558 | groups, g1, g2 = _mixed_page_groups() |
| 559 | # Overlay models resolve to a block-outer layout at backend selection (the |
| 560 | # model's backend declares it); mirror that here. |
| 561 | config = get_kv_cache_config_from_groups( |
| 562 | _mock_vllm_config("BLNHC"), groups, MEMORY |
| 563 | ) |
| 564 | assert config.num_blocks == MEMORY // _expected_bytes_per_block(groups) |
| 565 | assert _pool_bytes_per_block(groups) == _expected_bytes_per_block(groups) |
| 566 | |
| 567 | views = _bind(config, "BLNHC") |
| 568 | assert set(views) == set(g1) | set(g2) |
| 569 | assert views["swa.0"].data_ptr() == views["mla.0"].data_ptr() |
| 570 | |
| 571 | # A block is owned by one group at a time: writes to group-owned |
| 572 | # blocks never disturb the other group's blocks. |
| 573 | for i, name in enumerate(g1): |
| 574 | views[name][0].fill_(i + 1) |
| 575 | views[name][2].fill_(i + 1) |
| 576 | for i, name in enumerate(g2): |
| 577 | views[name][1].fill_(100 + i) |
| 578 | views[name][3].fill_(100 + i) |
| 579 | for i, name in enumerate(g1): |
| 580 | assert (views[name][0].to(torch.int32) == i + 1).all() |
| 581 | assert (views[name][2].to(torch.int32) == i + 1).all() |
| 582 | for i, name in enumerate(g2): |
| 583 | assert (views[name][1].to(torch.int32) == 100 + i).all() |
| 584 | assert (views[name][3].to(torch.int32) == 100 + i).all() |
| 585 | # Layers within a group are disjoint. |
| 586 | views["mla.0"][0].fill_(77) |
| 587 | for i, name in enumerate(list(g1)[1:], start=1): |
| 588 | assert (views[name][0].to(torch.int32) == i + 1).all() |
| 589 | |
| 590 | def test_layer_compact_layout_rejected_for_overlaid_groups(self): |
| 591 | # The layout has a single writer (backend-selection resolution); a layer- |
| 592 | # compact layout reaching an overlay model's allocation is an error, not a |
| 593 | # silent flip. |
| 594 | groups, _, _ = _mixed_page_groups() |
| 595 | with pytest.raises( |
| 596 | ValueError, match="cannot express this model's mixed page sizes" |
| 597 | ): |
| 598 | get_kv_cache_config_from_groups(_mock_vllm_config("LBNHC"), groups, MEMORY) |
| 599 | |
| 600 | def test_unresolved_layout_rejected(self): |
| 601 | groups, _, _ = _mixed_page_groups() |
| 602 | with pytest.raises(ValueError, match="has not been resolved"): |
| 603 | get_kv_cache_config_from_groups(_mock_vllm_config(None), groups, MEMORY) |
| 604 | |
| 605 | def test_head_outer_layout_rejected_for_mixed_pages(self): |
| 606 | groups, _, _ = _mixed_page_groups() |
| 607 | with pytest.raises( |
| 608 | ValueError, match="cannot express this model's mixed page sizes" |
| 609 | ): |
| 610 | get_kv_cache_config_from_groups(_mock_vllm_config("LHBNC"), groups, MEMORY) |
| 611 | |
| 612 | @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"]) |
| 613 | def test_bound_views_round_trip(self, layout): |
| 614 | specs = {"mla.0": _mla(512), "mla.1": _mla(512), "idx.0": _mla(128)} |
| 615 | groups = [_uniform_group(specs)] |
| 616 | config = get_kv_cache_config_from_groups( |
| 617 | _mock_vllm_config(layout), groups, MEMORY |
| 618 | ) |
| 619 | views = _bind(config, layout) |
| 620 | for i, name in enumerate(specs): |
| 621 | views[name].fill_(i + 1) |
| 622 | for i, name in enumerate(specs): |
| 623 | assert (views[name].to(torch.int32) == i + 1).all(), name |
| 624 | assert views[name].shape[0] == config.num_blocks |
| 625 | |
| 626 | |
| 627 | if __name__ == "__main__": |
| 628 | pytest.main([__file__, "-v"]) |
| 629 | |
| 630 | |
| 631 | class TestCompressorRingGroup: |
| 632 | def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self): |
| 633 | """DeepSeek-V4.1 compressor state: a CircularBufferSpec ring (capacity |
| 634 | 8) beside 128-token paged MLA groups and 64-token SWA groups. The ring |
| 635 | claims one block per request, never hashes, and does not disturb the |
| 636 | scheduler/hash block sizes the paged groups imply.""" |
| 637 | config = _shared_layout_config() |
| 638 | config.cache_config.enable_prefix_caching = True |
| 639 | specs: dict[str, KVCacheSpec] = { |
| 640 | "layers.2.attn": MLAAttentionSpec( |
| 641 | block_size=128, |
| 642 | num_kv_heads=1, |
| 643 | head_size=584, |
| 644 | dtype=torch.uint8, |
| 645 | tokens_per_state=2, |
| 646 | alignment=576, |
| 647 | ), |
| 648 | "layers.2.attn.compressor.state_cache": CircularBufferSpec( |
| 649 | block_size=8, |
| 650 | num_kv_heads=1, |
| 651 | head_size=1024, |
| 652 | head_size_v=0, |
| 653 | dtype=torch.float32, |
| 654 | ), |
| 655 | } |
| 656 | for layer in range(4): |
| 657 | specs[f"layers.{layer}.attn.swa_cache"] = SlidingWindowMLASpec( |
| 658 | block_size=64, |
| 659 | num_kv_heads=1, |
| 660 | head_size=584, |
| 661 | dtype=torch.uint8, |
| 662 | sliding_window=128, |
| 663 | alignment=576, |
| 664 | ) |
| 665 | groups = get_kv_cache_groups(config, specs) |
| 666 | ring_groups = [g for g in groups if "state_cache" in g.layer_names[0]] |
| 667 | assert len(ring_groups) == 1 |
| 668 | assert not ring_groups[0].kv_cache_spec.prefix_cacheable |
| 669 | kv_cache_config = get_kv_cache_config_from_groups( |
| 670 | config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups) |
| 671 | ) |
| 672 | assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (128, 64) |
| 673 | manager = KVCacheManager( |
| 674 | generate_scheduler_kv_cache_config([kv_cache_config]), |
| 675 | max_model_len=8192, |
| 676 | enable_caching=True, |
| 677 | hash_block_size=64, |
| 678 | scheduler_block_size=128, |
| 679 | ) |
| 680 | assert len(manager.coordinator.single_type_managers) == len(groups) |
| 681 | |
| 682 | |
| 683 | class TestSWABoundedReplayGrouping: |
| 684 | def test_non_cacheable_swa_leaves_hits_to_the_paged_group(self): |
| 685 | """Under SWA bounded replay only the paged MLA group hashes (hash |
| 686 | block = its block) and the coordinator probes one spec group.""" |
| 687 | config = _shared_layout_config() |
| 688 | config.cache_config.enable_prefix_caching = True |
| 689 | specs: dict[str, KVCacheSpec] = { |
| 690 | "layers.2.attn": MLAAttentionSpec( |
| 691 | block_size=128, |
| 692 | num_kv_heads=1, |
| 693 | head_size=584, |
| 694 | dtype=torch.uint8, |
| 695 | tokens_per_state=2, |
| 696 | alignment=576, |
| 697 | ), |
| 698 | } |
| 699 | swa = SlidingWindowMLASpec( |
| 700 | block_size=32, |
| 701 | num_kv_heads=1, |
| 702 | head_size=584, |
| 703 | dtype=torch.uint8, |
| 704 | sliding_window=128, |
| 705 | alignment=576, |
| 706 | bounded_replay=True, |
| 707 | ) |
| 708 | for layer in range(4): |
| 709 | specs[f"layers.{layer}.attn.swa_cache"] = swa |
| 710 | |
| 711 | groups = get_kv_cache_groups(config, specs) |
| 712 | # The worker reads the replay window off the (packed) group specs to |
| 713 | # arm the window clamp; the wrapper must forward it like cacheability. |
| 714 | assert sorted( |
| 715 | (g.kv_cache_spec.prefix_cacheable, g.kv_cache_spec.prefix_replay_tokens) |
| 716 | for g in groups |
| 717 | ) == [(False, 128)] * 4 + [(True, 0)] |
| 718 | kv_cache_config = get_kv_cache_config_from_groups( |
| 719 | config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups) |
| 720 | ) |
| 721 | manager = KVCacheManager( |
| 722 | generate_scheduler_kv_cache_config([kv_cache_config]), |
| 723 | max_model_len=8192, |
| 724 | enable_caching=True, |
| 725 | hash_block_size=128, |
| 726 | scheduler_block_size=128, |
| 727 | ) |
| 728 | assert len(manager.coordinator.attention_groups) == 1 |
| 729 | assert isinstance( |
| 730 | manager.coordinator.attention_groups[0].spec, MLAAttentionSpec |
| 731 | ) |
| 732 | |
| 733 | # Allocating and caching a request goes through every group's manager; |
| 734 | # the 32-token SWA group must skip hashing rather than assert on the |
| 735 | # 128-token hash block. |
| 736 | from tests.v1.core.utils import create_requests |
| 737 | |
| 738 | request = create_requests(num_requests=1, num_tokens=256, block_size=128)[0] |
| 739 | assert manager.allocate_slots(request, 256) is not None |
| 740 | request.num_computed_tokens = 256 |
| 741 | manager.cache_blocks(request, 256) |
| 742 | for single in manager.coordinator.single_type_managers: |
| 743 | cached = single.num_cached_block.get(request.request_id, 0) |
| 744 | assert cached == (2 if single.kv_cache_spec.prefix_cacheable else 0) |