@tduncan / Nvidia-Optimization-Stack-VLLM / commits / eb6b619ab2

[Bugfix][KV Offload] Register mixed page sizes in one cache group (#54756)

Signed-off-by: heliubj18 <chengfeng@ppio.com> Co-authored-by: Kimi <noreply@moonshot.cn>

He Liu committed Sep 8, 2026 at 11:50 UTC eb6b619ab256df1100cb75684a9ba78485506c3a
3 files changed +275 -31
tests/v1/attention/utils.py
+18 -6
@@ -426,24 +426,36 @@ class MockMambaBuilder(BaseMambaAttentionMetadataBuilder[BaseMambaAttentionMetad
426 return builder.build(0, common_metadata)
427
428
429 -def dense_kv_cache_views(
429 +def dense_kv_cache_tensor(
430 raw: torch.Tensor,
431 spec: KVCacheSpec,
432 num_blocks: int,
433 num_layers: int,
434 layout: KVCacheLayout,
435 - kernel_block_size: int | None = None,
436 -) -> list[torch.Tensor]:
437 - """``create_kv_cache_views`` for a dense allocation of ``num_layers`` layers."""
435 + layer_names: list[str] | None = None,
436 +) -> KVCacheTensor:
437 + """The ``KVCacheTensor`` for a dense allocation of ``num_layers`` layers."""
438 layer_stride, block_stride, _, _, _ = compute_layout_strides(
439 spec, num_blocks, num_layers, layout
440 )
441 - tensor = KVCacheTensor(
441 + return KVCacheTensor(
442 size=raw.numel(),
443 - layers=[str(i) for i in range(num_layers)],
443 + layers=layer_names or [str(i) for i in range(num_layers)],
444 layer_stride=layer_stride,
445 block_stride=block_stride,
446 )
447 +
448 +
449 +def dense_kv_cache_views(
450 + raw: torch.Tensor,
451 + spec: KVCacheSpec,
452 + num_blocks: int,
453 + num_layers: int,
454 + layout: KVCacheLayout,
455 + kernel_block_size: int | None = None,
456 +) -> list[torch.Tensor]:
457 + """``create_kv_cache_views`` for a dense allocation of ``num_layers`` layers."""
458 + tensor = dense_kv_cache_tensor(raw, spec, num_blocks, num_layers, layout)
459 return create_kv_cache_views(
460 raw,
461 spec,
tests/v1/simple_kv_offload/test_worker.py
+215 -5
@@ -20,8 +20,22 @@ from vllm.platforms import current_platform
20 if not current_platform.is_cuda_alike():
21 pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
22
23 -from tests.v1.attention.utils import dense_kv_cache_views
24 -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout
23 +from tests.v1.attention.utils import dense_kv_cache_tensor, dense_kv_cache_views
24 +from vllm.config import CacheConfig
25 +from vllm.v1.core.kv_cache_utils import (
26 + get_kv_cache_config_from_groups,
27 + is_kv_cache_spec_uniform,
28 + resolve_kv_cache_block_sizes,
29 +)
30 +from vllm.v1.kv_cache_interface import (
31 + FullAttentionSpec,
32 + KVCacheGroupSpec,
33 + KVCacheLayout,
34 + KVCacheSpec,
35 + KVCacheTensor,
36 + MLAAttentionSpec,
37 + UniformTypeKVCacheSpecs,
38 +)
39 from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend
40 from vllm.v1.simple_kv_offload.cuda_mem_ops import (
41 CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
@@ -31,6 +45,7 @@ from vllm.v1.simple_kv_offload.cuda_mem_ops import (
45 )
46 from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadMetadata
47 from vllm.v1.simple_kv_offload.worker import SimpleCPUOffloadWorker
48 +from vllm.v1.worker.utils import allocate_kv_cache
49
50 NUM_BLOCKS = 64
51 BLOCK_BYTES = 4096
@@ -210,9 +225,14 @@ def test_register_shared_kv_cache_storage(monkeypatch, layout: KVCacheLayout):
225 device="cuda",
226 )
227 caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
228 + layer_names = [f"layer.{i}" for i in range(num_layers)]
229 cache_config = MagicMock(
230 num_blocks=num_blocks,
215 - kv_cache_tensors=[MagicMock(size=raw.nbytes)],
231 + kv_cache_tensors=[
232 + dense_kv_cache_tensor(
233 + raw, spec, num_blocks, num_layers, layout, layer_names
234 + )
235 + ],
236 )
237 worker = SimpleCPUOffloadWorker(
238 vllm_config=None,
@@ -252,7 +272,14 @@ def test_register_kv_cache_storage_with_trailing_padding(monkeypatch):
272 vllm_config=None,
273 kv_cache_config=MagicMock(
274 num_blocks=num_blocks,
255 - kv_cache_tensors=[MagicMock(size=cache_bytes)],
275 + kv_cache_tensors=[
276 + KVCacheTensor(
277 + size=cache_bytes,
278 + layers=["layer.0"],
279 + layer_stride=cache_bytes,
280 + block_stride=block_bytes,
281 + )
282 + ],
283 ),
284 cpu_capacity_bytes=cache_bytes,
285 )
@@ -286,11 +313,16 @@ def test_register_separate_kv_head_groups(monkeypatch):
313 device="cuda",
314 )
315 caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
316 + layer_names = [f"layer.{i}" for i in range(num_layers)]
317 worker = SimpleCPUOffloadWorker(
318 vllm_config=None,
319 kv_cache_config=MagicMock(
320 num_blocks=num_blocks,
293 - kv_cache_tensors=[MagicMock(size=raw.nbytes)],
321 + kv_cache_tensors=[
322 + dense_kv_cache_tensor(
323 + raw, spec, num_blocks, num_layers, layout, layer_names
324 + )
325 + ],
326 ),
327 cpu_capacity_bytes=raw.nbytes,
328 )
@@ -309,3 +341,181 @@ def test_register_separate_kv_head_groups(monkeypatch):
341 assert {cache.shape for cache in worker.gpu_kv_caches.values()} == {
342 (num_blocks, per_group_block_bytes)
343 }
344 +
345 +
346 +def _dsa_specs(num_layers: int, block_size: int) -> dict[str, KVCacheSpec]:
347 + """A DSA model's per-layer specs (DeepSeek-V3.2, GLM-5.2).
348 +
349 + Each decoder layer contributes an MLA latent cache and the indexer's key
350 + cache, whose pages differ in size (``MLAAttentionSpec`` off ``head_size``:
351 + 512 + 64 for the latent, 128 fp8 keys + one fp32 scale per 128 elements for
352 + the indexer). Both are ``MLAAttentionSpec`` with the same block size, so the
353 + hybrid allocator puts them in one ``UniformTypeKVCacheSpecs`` cache group.
354 + See ``DeepseekV32IndexerCache.get_kv_cache_spec`` and
355 + ``MLAAttention.get_kv_cache_spec``.
356 + """
357 + specs: dict[str, KVCacheSpec] = {}
358 + for i in range(num_layers):
359 + specs[f"model.layers.{i}.self_attn.attn"] = MLAAttentionSpec(
360 + block_size=block_size,
361 + num_kv_heads=1,
362 + head_size=512 + 64,
363 + dtype=torch.uint8,
364 + cache_dtype_str="fp8",
365 + )
366 + specs[f"model.layers.{i}.self_attn.indexer.k_cache"] = MLAAttentionSpec(
367 + block_size=block_size,
368 + num_kv_heads=1,
369 + head_size=128 + 128 // 128 * 4,
370 + dtype=torch.uint8,
371 + )
372 + return specs
373 +
374 +
375 +def test_register_mixed_page_sizes_in_one_cache_group(monkeypatch):
376 + """Sparse-MLA models mix page sizes inside one cache group.
377 +
378 + The MLA latent and indexer key caches of a layer have different page sizes,
379 + so the allocation has no single per-layer block size to reinterpret it with:
380 + ``num_layers * (mla_page + indexer_page)`` is not a multiple of either page.
381 + Registration must derive each region from the placement metadata instead.
382 + """
383 + num_layers = 4
384 + block_size = 64
385 + specs = _dsa_specs(num_layers, block_size)
386 + # Differing head sizes make the specs non-identical but same-type, which is
387 + # what lands both caches of a layer in one group.
388 + assert not is_kv_cache_spec_uniform(specs)
389 + assert UniformTypeKVCacheSpecs.is_uniform_type(specs)
390 + group = KVCacheGroupSpec(
391 + list(specs),
392 + UniformTypeKVCacheSpecs(block_size=block_size, kv_cache_specs=specs),
393 + )
394 +
395 + # DEEPSEEK_V32_INDEXER declares no supported layouts, so resolution lands on
396 + # the default preference: layer-outermost, one region per layer.
397 + layout = KVCacheLayout.LBNHC
398 + vllm_config = MagicMock()
399 + vllm_config.cache_config = CacheConfig()
400 + vllm_config.cache_config.kv_cache_layout = layout.name
401 + vllm_config.cache_config.num_gpu_blocks_override = None
402 +
403 + pages = [spec.page_size_bytes for spec in specs.values()]
404 + num_blocks = 4
405 + kv_cache_config = get_kv_cache_config_from_groups(
406 + vllm_config, [group], sum(pages) * num_blocks
407 + )
408 + assert kv_cache_config.num_blocks == num_blocks
409 + assert len(set(pages)) > 1, "the mixed page sizes are what this test covers"
410 +
411 + kv_caches = allocate_kv_cache(kv_cache_config, torch.device("cuda"), layout)
412 + worker = SimpleCPUOffloadWorker(
413 + vllm_config=None,
414 + kv_cache_config=kv_cache_config,
415 + cpu_capacity_bytes=sum(pages) * num_blocks,
416 + )
417 + worker._backend = MagicMock()
418 + monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False)
419 +
420 + worker.register_kv_caches(kv_caches)
421 +
422 + assert worker.gpu_kv_caches is not None
423 + # One region per layer, each with that layer's own page as its block stride.
424 + assert sorted(worker.gpu_kv_caches) == sorted(specs)
425 + assert {name: cache.shape[1] for name, cache in worker.gpu_kv_caches.items()} == {
426 + name: spec.page_size_bytes for name, spec in specs.items()
427 + }
428 + assert all(cache.shape[0] == num_blocks for cache in worker.gpu_kv_caches.values())
429 +
430 + # Every registered region must alias the bytes the model writes through, or
431 + # offloaded blocks would be copied from the wrong place.
432 + for name, cache in kv_caches.items():
433 + region = worker.gpu_kv_caches[name]
434 + assert region.data_ptr() == cache.data_ptr()
435 + assert region.stride(0) == specs[name].page_size_bytes
436 +
437 +
438 +@pytest.mark.parametrize("rank_blocks", [1, 5])
439 +def test_register_mixed_page_sizes_odd_block_counts(monkeypatch, rank_blocks):
440 + """Registration holds at block counts that are not powers of two.
441 +
442 + ``num_blocks`` comes out of ``available_memory // bytes_per_block``, so
443 + any integer is reachable. Registration derives regions from the config's
444 + own block count and strides, and the CPU pool rounds down from the same
445 + per-block byte total; none of it assumes anything about how the count
446 + was reached.
447 + """
448 + num_layers = 4
449 + block_size = 64
450 + specs = _dsa_specs(num_layers, block_size)
451 + group = KVCacheGroupSpec(
452 + list(specs),
453 + UniformTypeKVCacheSpecs(block_size=block_size, kv_cache_specs=specs),
454 + )
455 + layout = KVCacheLayout.LBNHC
456 + vllm_config = MagicMock()
457 + vllm_config.cache_config = CacheConfig()
458 + vllm_config.cache_config.kv_cache_layout = layout.name
459 + vllm_config.cache_config.num_gpu_blocks_override = None
460 +
461 + pages = [spec.page_size_bytes for spec in specs.values()]
462 + kv_cache_config = get_kv_cache_config_from_groups(
463 + vllm_config, [group], sum(pages) * rank_blocks
464 + )
465 + assert kv_cache_config.num_blocks == rank_blocks
466 +
467 + kv_caches = allocate_kv_cache(kv_cache_config, torch.device("cuda"), layout)
468 + worker = SimpleCPUOffloadWorker(
469 + vllm_config=None,
470 + kv_cache_config=kv_cache_config,
471 + cpu_capacity_bytes=sum(pages) * rank_blocks,
472 + )
473 + worker._backend = MagicMock()
474 + monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False)
475 +
476 + worker.register_kv_caches(kv_caches)
477 +
478 + assert worker.gpu_kv_caches is not None
479 + assert {n: c.shape for n, c in worker.gpu_kv_caches.items()} == {
480 + name: (rank_blocks, spec.page_size_bytes) for name, spec in specs.items()
481 + }
482 + assert worker.num_cpu_blocks == rank_blocks
483 +
484 +
485 +def test_mixed_page_byte_placement_is_dcp_invariant():
486 + """DCP scales a block's token span, not its byte placement.
487 +
488 + ``get_kv_cache_config_from_groups`` never reads
489 + ``decode_context_parallel_size``, so every placement field registration
490 + derives regions from is identical at dcp=1 and dcp=2. What DCP scales is
491 + the block's token span, reported by ``resolve_kv_cache_block_sizes``;
492 + checking that it doubles proves DCP was actually in effect, or the
493 + identity above would be trivially true.
494 + """
495 + specs = _dsa_specs(4, block_size=128)
496 +
497 + placements = []
498 + for dcp_size in (1, 2):
499 + group = KVCacheGroupSpec(
500 + list(specs),
501 + UniformTypeKVCacheSpecs(block_size=128, kv_cache_specs=specs),
502 + )
503 + vllm_config = MagicMock()
504 + vllm_config.cache_config = CacheConfig()
505 + vllm_config.cache_config.block_size = 128
506 + vllm_config.cache_config.kv_cache_layout = KVCacheLayout.LBNHC.name
507 + vllm_config.cache_config.num_gpu_blocks_override = None
508 + vllm_config.parallel_config.decode_context_parallel_size = dcp_size
509 +
510 + page_bytes = sum(spec.page_size_bytes for spec in specs.values())
511 + config = get_kv_cache_config_from_groups(vllm_config, [group], page_bytes * 4)
512 + scheduler_block_size, _ = resolve_kv_cache_block_sizes(config, vllm_config)
513 + assert scheduler_block_size == 128 * dcp_size
514 + placements.append(
515 + [
516 + (t.size, tuple(t.layers), t.layer_stride, t.block_stride, t.offset)
517 + for t in config.kv_cache_tensors
518 + ]
519 + )
520 +
521 + assert placements[0] == placements[1]
vllm/v1/simple_kv_offload/worker.py
+42 -20
@@ -100,33 +100,55 @@ class SimpleCPUOffloadWorker:
100 assert self.kv_cache_config is not None
101 num_blocks = self.kv_cache_config.num_blocks
102 assert self.kv_cache_config.kv_cache_tensors
103 - logical_storage_bytes = self.kv_cache_config.kv_cache_tensors[0].size
103
105 - # The DMA backend copies whole blocks as base + block_id * stride(0),
106 - # so view each unique allocation as [num_blocks, block_bytes].
104 + # The DMA backend copies blocks as base + block_id * stride(0), so every
105 + # region is a [num_blocks, block_bytes] view. Block bytes come from each
106 + # tensor's placement -- layers in one allocation can differ in page size.
107 unique_gpu_caches: dict[str, torch.Tensor] = {}
108 seen: set[tuple[torch.device, int]] = set()
109 - for name, tensor in kv_caches.items():
110 - storage = tensor.untyped_storage()
111 - key = (tensor.device, storage.data_ptr())
112 - if key in seen:
109 + for kv_cache_tensor in self.kv_cache_config.kv_cache_tensors:
110 + name = next((n for n in kv_cache_tensor.layers if n in kv_caches), None)
111 + if name is None:
112 continue
114 - seen.add(key)
115 -
116 - physical_per_block, remainder = divmod(tensor.shape[0], num_blocks)
117 - assert remainder == 0, (
118 - f"KV cache {name!r} has {tensor.shape[0]} physical blocks, which "
119 - f"is not divisible by {num_blocks} scheduler blocks"
113 + num_layers = len(kv_cache_tensor.layers)
114 + block_bytes = kv_cache_tensor.block_stride
115 + blocks_span = num_blocks * block_bytes
116 + layer_outermost = kv_cache_tensor.layer_stride >= blocks_span
117 + if layer_outermost:
118 + # One region per layer, or per KV head group under LHBNC.
119 + start = kv_cache_tensor.offset
120 + span = num_layers * kv_cache_tensor.layer_stride
121 + layer_names = kv_cache_tensor.layers
122 + else:
123 + start = 0
124 + span = blocks_span
125 + layer_names = kv_cache_tensor.layers[:1]
126 +
127 + cache = kv_caches[name]
128 + raw = torch.empty(0, dtype=torch.int8, device=cache.device).set_(
129 + cache.untyped_storage()
130 )
121 - block_bytes = tensor.stride(0) * tensor.element_size() * physical_per_block
122 - raw = torch.empty(0, dtype=torch.int8, device=tensor.device).set_(storage)
123 - assert raw.numel() >= logical_storage_bytes, (
124 - f"KV cache {name!r} storage has {raw.numel()} bytes, smaller "
125 - f"than the configured {logical_storage_bytes}-byte allocation"
131 + assert raw.numel() >= start + span, (
132 + f"KV cache {name!r} storage has {raw.numel()} bytes, smaller than "
133 + f"the {span} bytes it places at offset {start}"
134 + )
135 + regions = raw[start : start + span].view(-1, num_blocks, block_bytes)
136 + groups_per_layer, remainder = divmod(len(regions), len(layer_names))
137 + assert remainder == 0, (
138 + f"KV cache {name!r} places {len(layer_names)} layers in "
139 + f"{len(regions)} block-strided regions"
140 )
127 - regions = raw[:logical_storage_bytes].view(-1, num_blocks, block_bytes)
141 for idx, region in enumerate(regions):
129 - key_name = name if len(regions) == 1 else f"{name}.{idx}"
142 + key = (region.device, region.data_ptr())
143 + if key in seen:
144 + continue
145 + seen.add(key)
146 + layer_name = layer_names[idx // groups_per_layer]
147 + key_name = (
148 + layer_name
149 + if groups_per_layer == 1
150 + else f"{layer_name}.{idx % groups_per_layer}"
151 + )
152 unique_gpu_caches[key_name] = region
153
154 # Compute per-tensor bytes_per_block. Tensors may have different