From 01192139bf022bec84e2cca3a3e36e8bb5293b5c Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Fri, 19 Jun 2026 12:55:42 -0400 Subject: [PATCH] [DSv4] Pack KV caches into contiguous per-block allocations for DeepSeek V4 (#44577) Signed-off-by: Tyler Michael Smith Signed-off-by: Matthew Bonanni Signed-off-by: Lucas Wilkinson Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Matthew Bonanni Co-authored-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- tests/v1/core/test_contiguous_kv_packing.py | 135 ++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 98 +++++++++++++ .../kv_connector/v1/offloading/worker.py | 14 +- vllm/v1/core/kv_cache_utils.py | 13 +- vllm/v1/kv_cache_interface.py | 2 + vllm/v1/worker/gpu/attn_utils.py | 48 ++++++- vllm/v1/worker/gpu_model_runner.py | 56 ++++++-- 7 files changed, 344 insertions(+), 22 deletions(-) create mode 100644 tests/v1/core/test_contiguous_kv_packing.py diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py new file mode 100644 index 0000000000..79f8937c63 --- /dev/null +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for contiguous KV cache packing in _get_kv_cache_config_deepseek_v4.""" + +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.v1.core.kv_cache_utils import _get_kv_cache_config_deepseek_v4 +from vllm.v1.kv_cache_interface import ( + KVCacheGroupSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, +) + + +def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec: + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=512, + dtype=torch.uint8, + page_size_padded=page_size, + cache_dtype_str="fp8_ds_mla", + model_version="deepseek_v4", + alignment=576, + ) + + +def _make_groups(n_c4, n_c128, n_swa): + PS_C4_MLA = 37440 + PS_C4_IDX = 8640 + PS_C128 = 1728 + PS_SWA = 37440 + + mla_specs = {} + for i in range(n_c4): + mla_specs[f"c4_mla.{i}"] = _make_mla_spec(PS_C4_MLA) + mla_specs[f"c4_idx.{i}"] = _make_mla_spec(PS_C4_IDX) + for i in range(n_c128): + mla_specs[f"c128_mla.{i}"] = _make_mla_spec(PS_C128) + + mla_group = KVCacheGroupSpec( + layer_names=list(mla_specs.keys()), + kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=mla_specs), + ) + + swa_specs = {} + for i in range(n_swa): + swa_specs[f"swa.{i}"] = _make_mla_spec(PS_SWA) + + swa_group = KVCacheGroupSpec( + layer_names=list(swa_specs.keys()), + kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=swa_specs), + ) + + return [mla_group, swa_group] + + +def _mock_vllm_config(): + config = MagicMock() + config.cache_config.num_gpu_blocks_override = None + return config + + +def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024): + groups = _make_groups(n_c4, n_c128, n_swa) + return _get_kv_cache_config_deepseek_v4(_mock_vllm_config(), groups, mem) + + +def _page_sizes_by_layer( + groups: list[KVCacheGroupSpec], +) -> dict[str, int]: + page_sizes = {} + for group in groups: + specs = group.kv_cache_spec.kv_cache_specs + for layer_name in group.layer_names: + page_sizes[layer_name] = specs[layer_name].page_size_bytes + return page_sizes + + +class TestInterleavedPacking: + def test_all_tensors_have_block_stride(self): + _, tensors = _run() + for t in tensors: + assert t.block_stride > 0 + + def test_all_tensors_share_same_size(self): + _, tensors = _run() + sizes = set(t.size for t in tensors) + assert len(sizes) == 1 + assert sizes.pop() > 0 + + def test_offsets_within_one_block(self): + _, tensors = _run() + for t in tensors: + assert t.offset < t.block_stride + + def test_all_layers_accounted_for(self): + n_c4, n_c128, n_swa = 5, 4, 7 + _, tensors = _run(n_c4=n_c4, n_c128=n_c128, n_swa=n_swa) + all_names = set() + for t in tensors: + all_names.update(t.shared_by) + expected = n_c4 * 2 + n_c128 + n_swa + assert len(all_names) == expected + + def test_strided_views_are_independent(self): + groups = _make_groups(n_c4=3, n_c128=2, n_swa=5) + page_sizes = _page_sizes_by_layer(groups) + num_blocks, tensors = _get_kv_cache_config_deepseek_v4( + _mock_vllm_config(), groups, 100 * 1024 * 1024 + ) + backing = torch.zeros(tensors[0].size, dtype=torch.uint8) + views = [] + for t in tensors: + page_size = page_sizes[t.shared_by[0]] + v = torch.as_strided( + backing, + size=(num_blocks, page_size), + stride=(t.block_stride, 1), + storage_offset=t.offset, + ) + views.append(v) + + for i, v in enumerate(views): + v.fill_(i + 1) + + for i, v in enumerate(views): + assert (v == i + 1).all(), f"View {i} was corrupted" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 5804732f80..7ee072ceaf 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -841,8 +841,106 @@ class NixlBaseConnectorWorker: # Forwarding a real layer name rather than a synthetic key self.register_kv_caches({first_layer: kv_cache}) + def _register_packed_kv_cache( + self, + storage: torch.UntypedStorage, + ) -> None: + """Register a packed KV cache as a single NIXL region. + + The packed allocation interleaves all layers per block, so each + block_stride-byte chunk is one logical block. We register 1 + NIXL region and create 1 descriptor per block. + """ + self.transfer_topo = TransferTopology( + tp_rank=self.tp_rank, + tp_size=self.world_size, + block_size=self.block_size, + engine_id=self.engine_id, + is_mla=self.use_mla, + total_num_kv_heads=self.model_config.get_total_num_kv_heads(), + attn_backends=self.attn_backends, + tensor_shape=None, + is_mamba=self._has_mamba, + ) + self.compat_hash = compute_nixl_compatibility_hash( + self.vllm_config, + self.backend_name, + self.transfer_topo.cross_layers_blocks, + ) + + total_size = storage.nbytes() + block_stride = total_size // self.num_blocks + base_addr = storage.data_ptr() + device_id = storage.device.index + assert device_id is not None + + logger.info( + "Registering packed KV cache: total_size=%s, block_stride=%s, " + "num_blocks=%s, num_regions=1", + total_size, + block_stride, + self.num_blocks, + ) + + self.device_id = device_id + caches_data = [(base_addr, total_size, self.device_id, "")] + + self.block_len_per_layer = [block_stride] + self.num_regions = 1 + self.num_descs = self.num_blocks + self.kv_caches_base_addr[self.engine_id][self.tp_rank] = [base_addr] + + descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type) + self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends) + self._registered_descs.append(descs) + + self.dst_num_blocks[self.engine_id] = self.num_blocks + + self.src_xfer_handles_by_block_size[self.block_size], (self.src_blocks_data) = ( + self.register_local_xfer_handler(self.block_size) + ) + + agent_metadata = NixlAgentMetadata( + engine_id=self.engine_id, + agent_metadata=self.nixl_wrapper.get_agent_metadata(), + device_id=self.device_id, + kv_caches_base_addr=( + self.kv_caches_base_addr[self.engine_id][self.tp_rank] + ), + num_blocks=self.num_blocks, + block_lens=self.block_len_per_layer, + kv_cache_layout=self.kv_cache_layout, + block_size=self.block_size, + ssm_sizes=self._mamba_ssm_size, + attn_backend_name=self.backend_name, + physical_blocks_per_logical_kv_block=( + self._physical_blocks_per_logical_kv_block + ), + ) + assert self.compat_hash is not None + encoder = msgspec.msgpack.Encoder() + self.xfer_handshake_metadata = NixlHandshakePayload( + compatibility_hash=self.compat_hash, + agent_metadata_bytes=encoder.encode(agent_metadata), + ) + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in nixl.""" + + # Detect packed allocation: all tensors are strided views into the + # same backing storage (different data_ptr but same storage). + # This happens with DSv4-style contiguous per-block packing. + if len(kv_caches) > 1 and not self._has_mamba: + storage = next(iter(kv_caches.values())).untyped_storage() + storage_ptrs = { + cache.untyped_storage().data_ptr() for cache in kv_caches.values() + } + data_ptrs = {cache.data_ptr() for cache in kv_caches.values()} + if len(storage_ptrs) == 1 and len(data_ptrs) > 1: + self._register_packed_kv_cache(storage) + self.device_kv_caches = kv_caches + return + self.transfer_topo = TransferTopology( tp_rank=self.tp_rank, tp_size=self.world_size, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 744a0c7429..8583bb4b1e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -72,18 +72,22 @@ class OffloadingConnectorWorker: if isinstance(layer_kv_cache_spec, AttentionSpec): layer_kv_cache = kv_caches[layer_name] assert isinstance(layer_kv_cache, torch.Tensor) - assert layer_kv_cache.storage_offset() == 0 - storage = layer_kv_cache.untyped_storage() page = layer_kv_cache_spec.page_size_bytes + elem_size = layer_kv_cache.element_size() + byte_offset = layer_kv_cache.storage_offset() * elem_size + block_stride_bytes = layer_kv_cache.stride(0) * elem_size tensors_per_block[layer_name] = ( torch.tensor( [], dtype=torch.int8, device=layer_kv_cache.device, - ) - .set_(storage) - .view(num_blocks, page), + ).set_( + layer_kv_cache.untyped_storage(), + byte_offset, + (num_blocks, page), + (block_stride_bytes, 1), + ), ) page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes unpadded_page_size_bytes[layer_name] = ( diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 72ca6a2fa6..a1ebe08c07 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1236,10 +1236,21 @@ def _get_kv_cache_config_deepseek_v4( num_blocks = available_memory // total_num_bytes_per_block num_blocks = may_override_num_blocks(vllm_config, num_blocks) + total_size = total_num_bytes_per_block * num_blocks + kv_cache_tensors: list[KVCacheTensor] = [] + byte_offset = 0 for ps, slots in buckets.items(): for slot in slots: - kv_cache_tensors.append(KVCacheTensor(size=ps * num_blocks, shared_by=slot)) + kv_cache_tensors.append( + KVCacheTensor( + size=total_size, + shared_by=slot, + offset=byte_offset, + block_stride=total_num_bytes_per_block, + ) + ) + byte_offset += ps return num_blocks, kv_cache_tensors diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 9528fb65af..2e779b2c2a 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -847,6 +847,8 @@ class KVCacheTensor: size: int # size of the KV cache tensor in bytes shared_by: list[str] # layer names that share the same KV cache tensor + offset: int = 0 # byte offset of this layer within a contiguous block + block_stride: int = 0 # total bytes per block in a packed layout (0 = not packed) @dataclass diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 74158f92bf..7b85e6fa31 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable, Sequence from dataclasses import dataclass +from math import prod from typing import Any, cast import torch @@ -155,8 +156,17 @@ def _allocate_kv_cache( kv_cache_config: KVCacheConfig, shared_layers: dict[str, str], device: torch.device ): kv_cache_raw_tensors: dict[str, torch.Tensor] = {} + packed_backing: torch.Tensor | None = None for kv_cache_tensor in kv_cache_config.kv_cache_tensors: - tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device) + if kv_cache_tensor.block_stride > 0: + # Allocate once; all packed tensors alias the same backing. + if packed_backing is None: + packed_backing = torch.zeros( + kv_cache_tensor.size, dtype=torch.int8, device=device + ) + tensor = packed_backing + else: + tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device) for layer_name in kv_cache_tensor.shared_by: kv_cache_raw_tensors[layer_name] = tensor @@ -176,10 +186,18 @@ def _reshape_kv_cache( cache_dtype: str, kernel_block_sizes: list[int], shared_kv_cache_layers: dict[str, str], + kv_cache_config: "KVCacheConfig | None" = None, ) -> dict[str, Any]: kv_caches: dict[str, Any] = {} has_attn, has_mamba = False, False + layer_packing: dict[str, tuple[int, int]] = {} + if kv_cache_config is not None: + for kv_tensor in kv_cache_config.kv_cache_tensors: + if kv_tensor.block_stride > 0: + for ln in kv_tensor.shared_by: + layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride) + for group in attn_groups: if group.kv_cache_group_id >= len(kernel_block_sizes): continue @@ -198,8 +216,13 @@ def _reshape_kv_cache( continue kv_raw_tensor = kv_cache_raw_tensors[layer_name] - assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 - num_blocks = kv_raw_tensor.numel() // kv_cache_spec.page_size_bytes + packing = layer_packing.get(layer_name) + if packing is not None: + _, blk_stride = packing + num_blocks = kv_raw_tensor.numel() // blk_stride + else: + assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 + num_blocks = kv_raw_tensor.numel() // kv_cache_spec.page_size_bytes if isinstance(kv_cache_spec, AttentionSpec): has_attn = True @@ -232,8 +255,18 @@ def _reshape_kv_cache( ] dtype = kv_cache_spec.dtype - kv_tensor = kv_raw_tensor.view(dtype) - if kv_cache_spec.page_size_padded is not None: + if packing is not None: + offset, block_stride = packing + assert inv_order[0] == 0 + page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype) + kv_cache = ( + kv_raw_tensor.view(-1, block_stride)[ + :, offset : offset + page_bytes + ] + .view(dtype) + .view(kv_cache_shape) + ) + elif kv_cache_spec.page_size_padded is not None: # Use strided view to handle page_size_bytes that # include padding. This follows the same pattern as # MambaSpec handling in gpu_model_runner.py. @@ -246,13 +279,13 @@ def _reshape_kv_cache( strides = list(torch.empty(kv_cache_shape).stride()) strides[inv_order[0]] = page_stride kv_cache = torch.as_strided( - kv_tensor, + kv_raw_tensor.view(dtype), size=kv_cache_shape, stride=tuple(strides), ) else: # No padding — safe to use a contiguous view. - kv_cache = kv_tensor.view(kv_cache_shape) + kv_cache = kv_raw_tensor.view(dtype).view(kv_cache_shape) kv_caches[layer_name] = kv_cache.permute(*inv_order) elif isinstance(kv_cache_spec, MambaSpec): @@ -365,6 +398,7 @@ def init_kv_cache( kernel_block_sizes=kernel_block_sizes, cache_dtype=cache_dtype, shared_kv_cache_layers=shared_kv_cache_layers, + kv_cache_config=kv_cache_config, ) bind_kv_cache(kv_caches, forward_context, runner_kv_caches) return kv_caches diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index b958ef79d0..3221dc46c6 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -12,6 +12,7 @@ from contextlib import contextmanager from copy import copy, deepcopy from dataclasses import dataclass, replace from functools import reduce +from math import prod from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast import numpy as np @@ -7029,10 +7030,21 @@ class GPUModelRunner( corresponding memory buffer for KV cache. """ kv_cache_raw_tensors: dict[str, torch.Tensor] = {} + packed_backing: torch.Tensor | None = None for kv_cache_tensor in kv_cache_config.kv_cache_tensors: - tensor = torch.zeros( - kv_cache_tensor.size, dtype=torch.int8, device=self.device - ) + if kv_cache_tensor.block_stride > 0: + # Allocate once; all packed tensors alias the same backing. + if packed_backing is None: + packed_backing = torch.zeros( + kv_cache_tensor.size, + dtype=torch.int8, + device=self.device, + ) + tensor = packed_backing + else: + tensor = torch.zeros( + kv_cache_tensor.size, dtype=torch.int8, device=self.device + ) for layer_name in kv_cache_tensor.shared_by: kv_cache_raw_tensors[layer_name] = tensor @@ -7074,6 +7086,14 @@ class GPUModelRunner( """ kv_caches: dict[str, torch.Tensor] = {} has_attn, has_mamba = False, False + + # Map layer names to (offset, block_stride) within the packed + # backing tensor so we can create strided views per layer. + layer_packing: dict[str, tuple[int, int]] = {} + for kv_tensor in self.kv_cache_config.kv_cache_tensors: + if kv_tensor.block_stride > 0: + for ln in kv_tensor.shared_by: + layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride) for group in self._kv_cache_spec_attn_group_iterator(): kv_cache_spec = group.kv_cache_spec attn_backend = group.backend @@ -7085,8 +7105,13 @@ class GPUModelRunner( if layer_name in self.runner_only_attn_layers: continue raw_tensor = kv_cache_raw_tensors[layer_name] - assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 - num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes + packing = layer_packing.get(layer_name) + if packing is not None: + _, blk_stride = packing + num_blocks = raw_tensor.numel() // blk_stride + else: + assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 + num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes if isinstance(kv_cache_spec, AttentionSpec): has_attn = True num_blocks_per_kv_block = ( @@ -7127,8 +7152,17 @@ class GPUModelRunner( for i in range(len(kv_cache_stride_order)) ] - raw_tensor = kv_cache_raw_tensors[layer_name].view(dtype) - if kv_cache_spec.page_size_padded is not None: + if packing is not None: + offset, block_stride = packing + assert inv_order[0] == 0 + page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype) + kv_cache = ( + kv_cache_raw_tensors[layer_name] + .view(-1, block_stride)[:, offset : offset + page_bytes] + .view(dtype) + .view(kv_cache_shape) + ) + elif kv_cache_spec.page_size_padded is not None: # Use strided view to handle page_size_bytes that # include padding. This follows # the same pattern as MambaSpec handling below. @@ -7142,13 +7176,17 @@ class GPUModelRunner( strides = list(torch.empty(kv_cache_shape).stride()) strides[inv_order[0]] = page_stride kv_cache = torch.as_strided( - raw_tensor, + kv_cache_raw_tensors[layer_name].view(dtype), size=kv_cache_shape, stride=tuple(strides), ) else: # No padding — safe to use a contiguous view. - kv_cache = raw_tensor.view(kv_cache_shape) + kv_cache = ( + kv_cache_raw_tensors[layer_name] + .view(dtype) + .view(kv_cache_shape) + ) kv_caches[layer_name] = kv_cache.permute(*inv_order) elif isinstance(kv_cache_spec, MambaSpec):