rebase: migrate diffusion KV cache to the vLLM 0.29 layout model

vLLM 0.29 removed AttentionSpec.indexes_kv_by_block_stride and reshaped
KVCacheTensor from (size, shared_by) to an explicit layout: layers in L
order plus layer_stride and block_stride. The block-stride decision now
lives on one physical CacheConfig.kv_cache_layout that the engine core
resolves once, and layout.is_block_outermost is the direct successor of
the old per-spec flag.

Diffusion backends subclass the Omni AttentionBackend ABC, so upstream
get_supported_kv_cache_layouts never sees them and the engine core never
resolves a layout for them. Add diffusion_kv/layout.py to pin it before
get_kv_cache_configs runs, defaulting to LBNHC, which reproduces the
indexes_kv_by_block_stride=False that every dense diffusion backend
declared. The backend hook is kept as an enforced contract rather than
deleted: a backend needing block-stride indexing selects a block-outermost
layout, and a contradiction raises instead of mis-addressing K/V.

Workers receive only a KVCacheConfig, so the resolved name travels on it
(KVCacheConfig.kv_cache_layout) and is adopted into each worker CacheConfig
before allocation. Without that, init_kv_cache raises in the worker because
get_resolved_kv_cache_layout has nothing to return.

init_kv_cache also lost attn_groups and cache_dtype in 0.29; the production
caller in model_runner_backend and the test stubs move to the six-argument
form.

Tensor construction goes through one derivation built on the upstream
compute_layout_strides helper, shared by production and tests, so a wrong
stride cannot be masked by a parallel test-only copy.

Also updates the sizing expectation: 0.29 subtracts the null block that
BlockPool permanently holds back before estimating, so sixteen physical
pages now yield fifteen usable ones.

One review point is addressed only in part, deliberately.
build_native_kv_cache_configs now takes indexes_kv_by_block_stride and
propagates a block-outermost layout end to end when it is set, with a test
covering control plane to worker adoption. Automatic discovery of that
requirement from the workers is not added: 0.29 removed the field that used
to carry it, every diffusion backend in tree declares False, and inventing
an executor RPC for a capability nothing requests would ship untestable
plumbing. Until a backend needs it, the seam is explicit and the mismatch
raises at worker init rather than corrupting addressing.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0146Xu5Sgj4vXxLxvLjV5YPj
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/diffusion/diffusion_kv/initialization.py

@@ -16,6 +16,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout
 from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
 from vllm_omni.platforms import current_omni_platform
 
@@ -28,14 +29,35 @@ def build_native_kv_cache_configs(
     vllm_config: VllmConfig,
     worker_specs: list[dict[str, KVCacheSpec]],
     available_memory: list[int],
+    *,
+    indexes_kv_by_block_stride: bool = False,
 ) -> tuple[list[KVCacheConfig], KVCacheConfig]:
-    """Build rank-local and Scheduler-native configs using vLLM utilities."""
+    """Build rank-local and Scheduler-native configs using vLLM utilities.
+
+    ``indexes_kv_by_block_stride`` carries the native backend requirement that
+    0.29 removed from ``AttentionSpec``.  Every diffusion backend in tree
+    declares ``False``, so the default reproduces current behaviour; a caller
+    that learns otherwise from its workers passes ``True`` here and the
+    block-outermost layout propagates to every rank-local config.
+    """
 
     from vllm.v1.core import single_type_kv_cache_manager
 
+    # 0.29 sizes every KVCacheTensor region from the resolved physical layout,
+    # so it must be pinned before get_kv_cache_configs.  Idempotent, so callers
+    # that already resolved it (the control plane) are unaffected.
+    layout = resolve_diffusion_kv_cache_layout(
+        vllm_config,
+        indexes_kv_by_block_stride=indexes_kv_by_block_stride,
+    )
     single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config)
     worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory)
-    return worker_configs, generate_scheduler_kv_cache_config(worker_configs)
+    scheduler_config = generate_scheduler_kv_cache_config(worker_configs)
+    # Workers only receive a KVCacheConfig, so the resolved layout travels on it
+    # and is adopted into each worker's own CacheConfig before allocation.
+    for config in (*worker_configs, scheduler_config):
+        config.kv_cache_layout = layout.name
+    return worker_configs, scheduler_config
 
 
 def initialize_diffusion_kv_control_plane(

vllm_omni/diffusion/diffusion_kv/layout.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Physical KV cache layout selection for the native Diffusion paged cache.
+
+vLLM 0.29 removed ``AttentionSpec.indexes_kv_by_block_stride`` and replaced it
+with a single physical layout per model, resolved once and recorded on
+``CacheConfig.kv_cache_layout``.  ``KVCacheLayout.is_block_outermost`` is the
+direct successor of the old per-spec flag: upstream derives
+``interleaved_block_stride`` from it when it lays out ``KVCacheTensor`` regions
+(``vllm/v1/core/kv_cache_utils.py``).
+
+Diffusion never reaches vLLM's engine-core resolution.  Its backends subclass
+Omni's own ``AttentionBackend`` ABC, which upstream's
+``get_supported_kv_cache_layouts`` never sees, so the layout is pinned here
+before ``get_kv_cache_configs`` runs.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from vllm.v1.kv_cache_interface import KVCacheSpec, KVCacheTensor, compute_layout_strides
+from vllm.v1.kv_cache_layout import KVCacheLayout
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+# Layer-outermost, matching the head of upstream's ``_DEFAULT_LAYOUT_PREFERENCE``
+# and reproducing the ``indexes_kv_by_block_stride=False`` that every dense
+# diffusion backend declared before 0.29.
+DEFAULT_DIFFUSION_KV_CACHE_LAYOUT = KVCacheLayout.LBNHC
+# Block-outermost counterpart, for a paged backend whose kernel really does read
+# K/V pages by the runtime block stride.
+BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT = KVCacheLayout.BLNHC
+
+
+def resolve_diffusion_kv_cache_layout(
+    vllm_config: VllmConfig,
+    *,
+    indexes_kv_by_block_stride: bool = False,
+) -> KVCacheLayout:
+    """Pin the physical KV layout for the Diffusion cache and return it.
+
+    A layout already present on the config wins, mirroring upstream's
+    ``resolve_kv_cache_layout``; it is validated against the requirement rather
+    than silently overwritten, because an inconsistent layout mis-addresses K/V
+    instead of raising.
+    """
+
+    required_block_outermost = bool(indexes_kv_by_block_stride)
+    cache_config = getattr(vllm_config, "cache_config", None)
+    if cache_config is None:
+        # A config double with nowhere to record the choice; report what the
+        # requirement implies without pretending it was pinned.
+        return (
+            BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT
+            if required_block_outermost
+            else DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+        )
+
+    current = getattr(cache_config, "kv_cache_layout", None)
+    if current is not None:
+        layout = KVCacheLayout[current]
+        if layout.is_block_outermost != required_block_outermost:
+            raise ValueError(
+                "Diffusion KV cache layout "
+                f"{layout.name} (is_block_outermost="
+                f"{layout.is_block_outermost}) contradicts the attention "
+                "backend, which requires is_block_outermost="
+                f"{required_block_outermost}."
+            )
+        return layout
+
+    layout = (
+        BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT
+        if required_block_outermost
+        else DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+    )
+    cache_config.kv_cache_layout = layout.name
+    return layout
+
+
+def adopt_kv_cache_layout(vllm_config: VllmConfig, kv_cache_config) -> KVCacheLayout:
+    """Adopt the control plane's resolved layout inside a worker.
+
+    The engine core resolves the layout once and every worker adopts it before
+    allocating (``CacheConfig.kv_cache_layout``).  Diffusion workers receive only
+    a ``KVCacheConfig``, so the name rides along on ``kv_cache_config`` and is
+    copied onto the worker's own config here.  Without this, anything calling
+    ``get_resolved_kv_cache_layout`` in a worker -- ``init_kv_cache`` does --
+    raises.
+    """
+
+    name = getattr(kv_cache_config, "kv_cache_layout", None)
+    if name is None:
+        # Older configs, or a standalone harness: fall back to the same default
+        # the control plane pins so both sides still agree.
+        return resolve_diffusion_kv_cache_layout(vllm_config)
+
+    cache_config = getattr(vllm_config, "cache_config", None)
+    if cache_config is None:
+        return KVCacheLayout[name]
+
+    current = getattr(cache_config, "kv_cache_layout", None)
+    if current is None:
+        cache_config.kv_cache_layout = name
+    elif current != name:
+        raise ValueError(
+            f"Worker KV cache layout {current!r} disagrees with the layout "
+            f"resolved by the control plane ({name!r})."
+        )
+    return KVCacheLayout[name]
+
+
+def build_kv_cache_tensor(
+    spec: KVCacheSpec,
+    num_blocks: int,
+    layers: Sequence[str],
+    *,
+    layout: KVCacheLayout | None = None,
+    offset: int = 0,
+) -> KVCacheTensor:
+    """A ``KVCacheTensor`` covering ``layers`` for ``num_blocks`` blocks.
+
+    Before 0.29 a tensor named the layers sharing it
+    (``KVCacheTensor(size=..., shared_by=[...])``); 0.29 replaced that with an
+    explicit layout, so layer ``l``'s block ``b`` starts at
+    ``offset + l * layer_stride + b * block_stride``.  Note this lays the layers
+    out at distinct offsets rather than aliasing them, which is what upstream
+    now does for a single cache group.
+
+    The strides come from upstream's own ``compute_layout_strides``, called the
+    way ``vllm.v1.core.kv_cache_utils`` calls it, because hand-computed strides
+    mis-address K/V silently instead of raising.
+    """
+
+    layout = layout or DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+    layer_names = list(layers)
+    # Mirrors kv_cache_utils: the block stride is pinned up front only for
+    # block-outermost layouts, where every block repeats the same packing.
+    bytes_per_block = spec.page_size_bytes * len(layer_names)
+    interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None
+    layer_stride, block_stride, _, _, _ = compute_layout_strides(
+        spec,
+        num_blocks,
+        len(layer_names),
+        layout,
+        fixed_strides=(None, interleaved_block_stride, None, None, None),
+    )
+    return KVCacheTensor(
+        size=bytes_per_block * num_blocks,
+        layers=layer_names,
+        layer_stride=layer_stride,
+        block_stride=block_stride,
+        offset=offset,
+    )
+
+
+def assert_backend_layout_supported(vllm_config: VllmConfig, attn_backend: type | None) -> None:
+    """Fail loudly when a backend's block-stride need contradicts the layout.
+
+    Skipped while the layout is still unresolved -- worker processes collect
+    specs before the control plane hands them the resolved name -- so this
+    enforces the contract wherever it is knowable and never crashes a rank that
+    has simply not been told yet.
+    """
+
+    if attn_backend is None:
+        return
+    # Spec discovery runs against lightweight config doubles in tests and against
+    # worker configs that have not been handed the resolved name yet; in both
+    # cases there is simply nothing to check.
+    cache_config = getattr(vllm_config, "cache_config", None)
+    if cache_config is None or getattr(cache_config, "kv_cache_layout", None) is None:
+        return
+    indexes_by_block_stride = getattr(attn_backend, "indexes_kv_by_block_stride", None)
+    if indexes_by_block_stride is None:
+        return
+    resolve_diffusion_kv_cache_layout(
+        vllm_config,
+        indexes_kv_by_block_stride=indexes_by_block_stride(),
+    )

vllm_omni/diffusion/diffusion_kv/model_runner_backend.py

@@ -18,6 +18,10 @@ from vllm.v1.worker.gpu.block_table import BlockTables
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode
+from vllm_omni.diffusion.diffusion_kv.layout import (
+    adopt_kv_cache_layout,
+    assert_backend_layout_supported,
+)
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionAdapter,
@@ -285,13 +289,18 @@ class DiffusionKVModelRunnerBackend:
                     layers=self._kv_cache_layer_adapters,
                     resolve_row=self._resolve_paged_attention_row,
                 )
+                # vLLM 0.29 reads the resolved physical layout when allocating,
+                # and dropped attn_groups/cache_dtype from the signature.
+                adopt_kv_cache_layout(self.vllm_config, kv_cache_config)
+                for layer_adapter in self._kv_cache_layer_adapters.values():
+                    assert_backend_layout_supported(
+                        self.vllm_config, getattr(layer_adapter, "attn_backend", None)
+                    )
                 init_kv_cache(
                     kv_caches,
                     self.vllm_config.compilation_config.static_forward_context,
                     kv_cache_config,
-                    attn_groups,
                     self.device,
-                    self.vllm_config.cache_config.cache_dtype,
                     kernel_block_sizes,
                     self.vllm_config,
                 )

vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py

@@ -105,10 +105,12 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
             attn_backend,
             ulysses_degree=ulysses_degree,
         )
+        # vLLM 0.29 carries the block-stride decision on the resolved physical
+        # ``CacheConfig.kv_cache_layout`` instead of the spec; the backend's
+        # requirement is enforced against that layout rather than copied here.
         canonical_spec = replace(
             spec,
             num_kv_heads=num_kv_heads,
-            indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(),
         )
         self.layer_name = layer_name
         self.spec = canonical_spec