Simplify KV cache layout stride handling

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Author: LucasWilkinson

tests/v1/attention/utils.py

@@ -35,8 +35,8 @@
     KVCacheLayout,
     KVCacheSpec,
     MambaSpec,
+    compute_layout_strides,
     get_kv_quant_mode,
-    layer_kv_cache_strides,
     reshape_kv_cache,
 )
 
@@ -432,7 +432,7 @@ def dense_kv_cache_views(
     block_size: int | None = None,
 ) -> list[torch.Tensor]:
     """``reshape_kv_cache`` for a dense allocation of ``num_layers`` layers."""
-    layer_stride, block_stride = layer_kv_cache_strides(
+    layer_stride, block_stride, _, _, _ = compute_layout_strides(
         spec, num_blocks, num_layers, layout, block_size
     )
     return reshape_kv_cache(

vllm/config/cache.py

@@ -57,10 +57,8 @@ class CacheConfig:
     kv_cache_layout: str | None = field(default=None, init=False)
     """Resolved physical KV cache layout name (a ``KVCacheLayout`` member).
 
-    Written exactly once by attention backend selection (priority:
-    test override > backend-required > VLLM_KV_CACHE_LAYOUT > connector
-    preference > LBNHC); every consumer reads it from here so allocation,
-    validation, and connectors can never disagree. Derived automatically."""
+    Resolved by attention backend selection and synchronized across engine and
+    worker processes before allocation. Derived automatically."""
     prefix_match_unit: int | None = Field(default=None, gt=0)
     """The finest token boundary (in tokens) a prefix-cache hit can land on.
 

vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py

@@ -34,10 +34,16 @@
 def canonical_format_id() -> str:
     """Identity of the canonical byte format, for namespacing persisted KV.
     Canonical pages keep the worker's KV layout family, so the id couples the
-    format version with that family; consumers must match it exactly."""
+    format version with that family; consumers must match it exactly.
+    The family keeps its historical NHD/HND spelling so ids stay stable for
+    KV persisted before the layout enum existed."""
     from vllm.v1.attention.backends.utils import get_kv_cache_layout
+    from vllm.v1.kv_cache_interface import KVCacheLayout
 
-    return f"v{CANONICAL_FORMAT_VERSION}-{get_kv_cache_layout().name.lower()}"
+    layout = get_kv_cache_layout()
+    legacy = {KVCacheLayout.LBNHC: "nhd", KVCacheLayout.LBHNC: "hnd"}
+    family = legacy.get(layout, layout.name.lower())
+    return f"v{CANONICAL_FORMAT_VERSION}-{family}"
 
 
 @dataclass(frozen=True)

vllm/v1/attention/backends/utils.py

@@ -214,12 +214,11 @@ def _validate_backend_supports_layout(
 def initialize_kv_cache_layout(
     backend: type[AttentionBackend], cache_config=None
 ) -> None:
-    """Resolve the layout once at backend selection and publish it.
+    """Resolve the backend-compatible layout and publish it on the worker.
 
-    Single writer for the resolved layout: stores it on
-    ``cache_config.kv_cache_layout`` (serialized with the config) and in a
-    process-local mirror for callers without a config handle. Priority is
-    main-parity: a backend-required layout silently corrects the env var.
+    Stores it on ``cache_config.kv_cache_layout`` and in a process-local mirror
+    for callers without a config handle. A backend-required layout silently
+    corrects the environment setting.
     """
     global _RESOLVED_KV_CACHE_LAYOUT, _RESOLVED_LAYOUT_REQUIRED_BY
     if _KV_CACHE_LAYOUT_OVERRIDE is not None:
@@ -271,7 +270,7 @@ def require_block_outer_kv_cache_layout(cache_config=None) -> KVCacheLayout:
     honored and raises.
     """
     global _RESOLVED_KV_CACHE_LAYOUT
-    layout = get_kv_cache_layout()
+    layout = get_kv_cache_layout(cache_config)
     if not layout.is_layer_compact:
         return layout
 
@@ -290,18 +289,18 @@ def require_block_outer_kv_cache_layout(cache_config=None) -> KVCacheLayout:
     return layout
 
 
-def get_kv_cache_layout() -> KVCacheLayout:
+def get_kv_cache_layout(cache_config=None) -> KVCacheLayout:
     """Return the resolved physical KV cache layout.
 
-    Read-only: prefers the test override, then the value published by
-    ``initialize_kv_cache_layout`` (via the process mirror or the current
-    vllm config), then falls back to env > connector > LBNHC for processes
-    where backend selection never runs (e.g. the engine core).
+    Read-only: prefers the test override, then an explicit or current config,
+    then the process-local value published by ``initialize_kv_cache_layout``.
+    Processes where backend selection never runs fall back to
+    env > connector > LBNHC.
     """
     if _KV_CACHE_LAYOUT_OVERRIDE is not None:
         return _layout_from_name(_KV_CACHE_LAYOUT_OVERRIDE)
-    if _RESOLVED_KV_CACHE_LAYOUT is not None:
-        return _RESOLVED_KV_CACHE_LAYOUT
+    if cache_config is not None and cache_config.kv_cache_layout is not None:
+        return _layout_from_name(cache_config.kv_cache_layout)
 
     from vllm.config import get_current_vllm_config_or_none
 
@@ -312,6 +311,8 @@ def get_kv_cache_layout() -> KVCacheLayout:
         and vllm_config.cache_config.kv_cache_layout is not None
     ):
         return _layout_from_name(vllm_config.cache_config.kv_cache_layout)
+    if _RESOLVED_KV_CACHE_LAYOUT is not None:
+        return _RESOLVED_KV_CACHE_LAYOUT
 
     layout_name = envs.VLLM_KV_CACHE_LAYOUT
     if layout_name is None:

vllm/v1/core/kv_cache_utils.py

@@ -38,7 +38,7 @@
     SlidingWindowMLASpec,
     SlidingWindowSpec,
     UniformTypeKVCacheSpecs,
-    layer_kv_cache_strides,
+    compute_layout_strides,
     replace_as,
 )
 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
@@ -1270,7 +1270,7 @@ def _resolve_layout_for_groups(
     cache_config=None,
 ) -> KVCacheLayout:
     """Resolve the layout this model's packing can be expressed in."""
-    layout = get_kv_cache_layout()
+    layout = get_kv_cache_layout(cache_config)
     page_sizes = {
         _get_per_layer_spec(group, layer_name).page_size_bytes
         for group in kv_cache_groups
@@ -1324,28 +1324,27 @@ def get_kv_cache_config_from_groups(
 
     layout = _resolve_layout_for_groups(kv_cache_groups, vllm_config.cache_config)
     kv_cache_tensors = []
-    dense = len(runs) == 1
     for byte_offset, layer_names, spec in runs:
-        page = spec.page_size_bytes
-        if dense:
-            layer_stride, block_stride = layer_kv_cache_strides(
-                spec, num_blocks, len(layer_names), layout
-            )
-        elif layout.is_layer_compact:
-            layer_stride, block_stride = page * num_blocks, page
-        else:
-            layer_stride, block_stride = page, packed_block_stride
+        strides = compute_layout_strides(
+            spec,
+            num_blocks,
+            len(layer_names),
+            layout,
+            packed_block_stride=packed_block_stride,
+        )
+        layer_stride, block_stride, _, _, _ = strides  # L, B, H, N, C
+        offset = (
+            byte_offset
+            * max(layer_stride, spec.page_size_bytes)
+            // spec.page_size_bytes
+        )
         kv_cache_tensors.append(
             KVCacheTensor(
                 size=size,
                 layers=layer_names,
                 layer_stride=layer_stride,
                 block_stride=block_stride,
-                # Layer-outermost layouts give each layer its own region, so
-                # a run starts past every preceding run's whole region.
-                offset=byte_offset * num_blocks
-                if layout.is_layer_compact
-                else byte_offset,
+                offset=offset,
             )
         )
 

vllm/v1/executor/abstract.py

@@ -117,7 +117,10 @@ def _init_executor(self) -> None:
 
     def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
         """Initialize the KV caches on the underlying workers."""
-        self.collective_rpc("initialize_from_config", args=(kv_cache_configs,))
+        self.collective_rpc(
+            "initialize_from_config",
+            args=(kv_cache_configs, self.cache_config.kv_cache_layout),
+        )
 
     def compile_or_warm_up_model(self) -> None:
         """Compile/warm up the model and capture cudagraphs on workers."""
@@ -147,7 +150,14 @@ def determine_available_memory(self) -> list[int]:  # in bytes
         return self.collective_rpc("determine_available_memory")
 
     def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
-        return self.collective_rpc("get_kv_cache_spec")
+        specs: list[dict[str, KVCacheSpec]] = self.collective_rpc("get_kv_cache_spec")
+        worker_layouts: list[str | None] = self.collective_rpc("get_kv_cache_layout")
+        layouts = {layout for layout in worker_layouts if layout is not None}
+        if len(layouts) > 1:
+            raise ValueError(f"Workers disagree on KV cache layout: {layouts}")
+        if layouts:
+            self.cache_config.kv_cache_layout = next(iter(layouts))
+        return specs
 
     @overload
     def collective_rpc(

vllm/v1/kv_cache_interface.py

@@ -280,39 +280,41 @@ def compute_layer_kv_cache_shape_bytes(
     return (num_blocks, spec.num_heads, num_states, spec.state_content_size_bytes)
 
 
-def layer_kv_cache_strides(
+def compute_layout_strides(
     spec: KVCacheSpec,
     num_blocks: int,
     num_layers: int,
     layout: KVCacheLayout,
     block_size: int | None = None,
-) -> tuple[int, int]:
-    """Byte ``(layer_stride, block_stride)`` of a dense ``[L, B, H, N, C]``
-    allocation in ``layout`` order."""
+    packed_block_stride: int | None = None,
+) -> tuple[int, ...]:
+    """Byte strides in logical ``[L, B, H, N, C]`` axis order."""
     shape = (
         num_layers,
         *compute_layer_kv_cache_shape_bytes(spec, num_blocks, block_size),
     )
     stride_order = layout.stride_order
     physical_shape = tuple(shape[i] for i in stride_order)
-    dense = torch.empty(physical_shape, device="meta").stride()
     inv_order = [stride_order.index(i) for i in range(5)]
-    layer_stride = dense[inv_order[_DIM_L]]
-    block_stride = dense[inv_order[_DIM_B]]
 
-    if padded := getattr(spec, "page_size_padded", None):
+    padded = getattr(spec, "page_size_padded", None)
+    if padded is not None:
         assert block_size is None or block_size == spec.block_size, (
             "Padded KV pages do not support kernel block splitting."
         )
         assert {inv_order[_DIM_L], inv_order[_DIM_B]} == {0, 1}, (
             f"Padded KV pages need L and B outermost, got {layout.name}."
         )
-        # Padding widens every page, so the strides that step over whole
-        # pages scale with it.
-        page = prod(shape[2:])
-        layer_stride = layer_stride // page * padded
-        block_stride = block_stride // page * padded
-    return layer_stride, block_stride
+
+    logical_tail = prod(physical_shape[2:])
+    storage_tail = padded if padded is not None else logical_tail
+    physical = torch.empty((*physical_shape[:2], storage_tail), device="meta")
+    strides = list(
+        physical[..., :logical_tail].view(physical_shape).permute(*inv_order).stride()
+    )
+    if packed_block_stride is not None and inv_order[_DIM_B] < inv_order[_DIM_L]:
+        strides[_DIM_B] = packed_block_stride
+    return tuple(strides)
 
 
 def reshape_kv_cache(
@@ -338,21 +340,21 @@ def reshape_kv_cache(
     # e.g. BHLNC's head stride spans the layers it interleaves); the caller's
     # strides place the layers and blocks themselves.
     logical_shape = (num_layers, *shape_bytes)
-    stride_order = layout.stride_order
-    physical_shape = tuple(logical_shape[i] for i in stride_order)
-    inv_order = [stride_order.index(i) for i in range(5)]
-    strides = list(torch.empty(physical_shape, device="meta").stride())
-    strides[inv_order[_DIM_L]] = layer_stride
-    strides[inv_order[_DIM_B]] = block_stride
+    strides = list(
+        compute_layout_strides(
+            spec, num_blocks, num_layers, layout, block_size=block_size
+        )
+    )
+    strides[_DIM_L] = layer_stride
+    strides[_DIM_B] = block_stride
     dtype = getattr(spec, "dtype", None)
 
-    cache = torch.as_strided(
+    cache_logical_5d = torch.as_strided(
         raw,
-        size=physical_shape,
+        size=logical_shape,
         stride=tuple(strides),
         storage_offset=raw.storage_offset() + offset,
     )
-    cache_logical_5d = cache.permute(*inv_order)
 
     views = []
     for layer_idx in range(num_layers):

vllm/v1/worker/worker_base.py

@@ -99,6 +99,11 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
         """Get specifications for KV cache implementation."""
         raise NotImplementedError
 
+    def get_kv_cache_layout(self) -> str | None:
+        """Return the worker's resolved KV cache layout."""
+        cache_config = getattr(self.vllm_config, "cache_config", None)
+        return getattr(cache_config, "kv_cache_layout", None)
+
     def compile_or_warm_up_model(self) -> CompilationTimes:
         """Prepare model for execution through compilation/warmup.
 
@@ -322,9 +327,15 @@ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None:
             # To make vLLM config available during worker initialization
             self.worker = worker_class(**kwargs)
 
-    def initialize_from_config(self, kv_cache_configs: list[Any]) -> None:
+    def initialize_from_config(
+        self,
+        kv_cache_configs: list[Any],
+        kv_cache_layout: str | None = None,
+    ) -> None:
         kv_cache_config = kv_cache_configs[self.global_rank]
         assert self.vllm_config is not None
+        if kv_cache_layout is not None:
+            self.vllm_config.cache_config.kv_cache_layout = kv_cache_layout
         with set_current_vllm_config(self.vllm_config):
             self.worker.initialize_from_config(kv_cache_config)  # type: ignore