[DSv4] Pack KV caches into contiguous per-block allocations for DeepSeek V4 (#44577)

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

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"])

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,

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] = (

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
 

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

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

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):