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>

tests/diffusion/diffusion_kv/test_block_tables.py

@@ -279,17 +279,17 @@ def _install_native_modules(monkeypatch: pytest.MonkeyPatch, events: list[tuple]
         group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None))
         return [[group], [group]], "cg-support", [4, 8]
 
+    # vLLM 0.29 dropped attn_groups and cache_dtype from init_kv_cache: the
+    # layout now rides on the resolved CacheConfig.kv_cache_layout instead.
     def init_cache(
         runner_kv_caches,
         forward_context,
         kv_cache_config,
-        attn_groups,
         device,
-        cache_dtype,
         kernel_block_sizes,
         vllm_config,
     ):
-        events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes))
+        events.append(("cache", kv_cache_config, kernel_block_sizes))
         runner_kv_caches.extend(["cache-0", "cache-1"])
         return {"layer-0": "tensor-0", "layer-1": "tensor-1"}
 

tests/diffusion/diffusion_kv/test_initialization.py

@@ -299,8 +299,10 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None
     )
 
     assert vllm_config.model_config.original_max_model_len == -1
-    # Sixteen physical pages provide ``16 * block_size`` tokens.  The native
-    # backend chooses the block geometry, so keep this assertion backend
+    # Sixteen physical pages provide ``(16 - 1) * block_size`` tokens: vLLM 0.29
+    # subtracts the null block that BlockPool permanently holds back before
+    # sizing (check_enough_kv_cache_memory), so only 15 pages are usable.  The
+    # native backend chooses the block geometry, so keep this assertion backend
     # independent instead of assuming the generic 16-token block size.
-    expected_max_model_len = min(model_limit, spec.block_size * 16)
+    expected_max_model_len = min(model_limit, spec.block_size * (16 - 1))
     assert vllm_config.model_config.max_model_len == expected_max_model_len

tests/diffusion/diffusion_kv/test_layout.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Physical KV cache layout selection and propagation for the Diffusion cache."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from vllm.config import CacheConfig
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+from vllm.v1.kv_cache_layout import KVCacheLayout
+
+from vllm_omni.diffusion.diffusion_kv.layout import (
+    DEFAULT_DIFFUSION_KV_CACHE_LAYOUT,
+    adopt_kv_cache_layout,
+    assert_backend_layout_supported,
+    build_kv_cache_tensor,
+    resolve_diffusion_kv_cache_layout,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _config():
+    return SimpleNamespace(cache_config=CacheConfig())
+
+
+class _DenseBackend:
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        return False
+
+
+class _BlockStrideBackend:
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        return True
+
+
+def test_default_layout_is_not_block_outermost() -> None:
+    """The default reproduces the pre-0.29 indexes_kv_by_block_stride=False."""
+    vllm_config = _config()
+    layout = resolve_diffusion_kv_cache_layout(vllm_config)
+
+    assert layout is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+    assert layout.is_block_outermost is False
+    assert vllm_config.cache_config.get_resolved_kv_cache_layout() is layout
+
+
+def test_block_stride_backend_selects_block_outermost_layout() -> None:
+    layout = resolve_diffusion_kv_cache_layout(_config(), indexes_kv_by_block_stride=True)
+
+    assert layout.is_block_outermost is True
+
+
+def test_resolution_is_idempotent() -> None:
+    vllm_config = _config()
+
+    assert resolve_diffusion_kv_cache_layout(vllm_config) is resolve_diffusion_kv_cache_layout(vllm_config)
+
+
+def test_contradicting_requirement_raises_instead_of_misaddressing() -> None:
+    vllm_config = _config()
+    resolve_diffusion_kv_cache_layout(vllm_config)
+
+    with pytest.raises(ValueError, match="contradicts the attention"):
+        resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True)
+
+
+def test_preset_layout_is_honoured() -> None:
+    vllm_config = _config()
+    vllm_config.cache_config.kv_cache_layout = KVCacheLayout.BLNHC.name
+
+    layout = resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True)
+
+    assert layout is KVCacheLayout.BLNHC
+
+
+def test_backend_assertion_is_quiet_until_the_layout_is_resolved() -> None:
+    """Workers collect specs before the control plane hands them the layout."""
+    vllm_config = _config()
+
+    assert_backend_layout_supported(vllm_config, _BlockStrideBackend)  # unresolved: nothing to check
+
+    resolve_diffusion_kv_cache_layout(vllm_config)
+    assert_backend_layout_supported(vllm_config, _DenseBackend)
+    assert_backend_layout_supported(vllm_config, None)
+    with pytest.raises(ValueError, match="contradicts the attention"):
+        assert_backend_layout_supported(vllm_config, _BlockStrideBackend)
+
+
+def test_worker_adopts_the_layout_carried_on_the_kv_cache_config() -> None:
+    """Diffusion workers receive only a KVCacheConfig, so the name rides on it."""
+    worker = _config()
+    shipped = SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name)
+
+    assert adopt_kv_cache_layout(worker, shipped) is KVCacheLayout.BLNHC
+    assert worker.cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.BLNHC
+
+
+def test_worker_without_a_carried_layout_falls_back_to_the_same_default() -> None:
+    worker = _config()
+
+    assert adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=None)) is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+
+
+def test_worker_disagreeing_with_the_control_plane_raises() -> None:
+    worker = _config()
+    resolve_diffusion_kv_cache_layout(worker)
+
+    with pytest.raises(ValueError, match="disagrees with the layout"):
+        adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name))
+
+
+def test_block_stride_requirement_propagates_control_plane_to_worker() -> None:
+    """A block-stride backend must reach the worker as a block-outermost layout.
+
+    Covers the whole seam: the control plane resolves with the requirement, the
+    name rides on each rank-local KVCacheConfig, the worker adopts it, and the
+    backend contract check then passes instead of raising.
+    """
+    control_plane = _config()
+    layout = resolve_diffusion_kv_cache_layout(control_plane, indexes_kv_by_block_stride=True)
+    assert layout.is_block_outermost is True
+
+    shipped = SimpleNamespace(kv_cache_layout=layout.name)
+
+    worker = _config()
+    assert adopt_kv_cache_layout(worker, shipped) is layout
+    # The worker's own backend requirement now agrees with what it was handed.
+    assert_backend_layout_supported(worker, _BlockStrideBackend)
+    with pytest.raises(ValueError, match="contradicts the attention"):
+        assert_backend_layout_supported(worker, _DenseBackend)
+
+
+def test_tensor_strides_come_from_the_upstream_derivation() -> None:
+    spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=8, dtype=torch.bfloat16)
+
+    tensor = build_kv_cache_tensor(spec, 4, ["layer0", "layer1"])
+
+    assert tensor.layers == ["layer0", "layer1"]
+    assert tensor.size == spec.page_size_bytes * 2 * 4
+    # Layer-outermost: each layer owns a contiguous run of its own blocks.
+    assert tensor.layer_stride == spec.page_size_bytes * 4
+    assert tensor.block_stride == spec.page_size_bytes

tests/diffusion/diffusion_kv/test_manager.py

@@ -4,10 +4,11 @@
 import pytest
 import torch
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
 from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
+from tests.helpers.kv_layout import build_kv_cache_tensor
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -24,7 +25,7 @@ def _config(num_blocks: int) -> KVCacheConfig:
     )
     return KVCacheConfig(
         num_blocks=num_blocks,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
     )
 

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -18,6 +18,7 @@ from vllm.v1.kv_cache_interface import FullAttentionSpec
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module
+from tests.helpers.kv_layout import layout_for_backend
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionAdapter,
     DiffusionPagedAttentionMetadata,
@@ -1239,7 +1240,8 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
     assert native_layer.num_heads == 4
     assert native_layer.num_kv_heads == 2
     assert native_layer.spec.num_kv_heads == 2
-    assert native_layer.spec.indexes_kv_by_block_stride is True
+    # 0.29: the flag no longer rides on the spec, it selects the layout.
+    assert layout_for_backend(native_layer.attn_backend).is_block_outermost is True
     assert native_layer._q_scale_float == 1.0
     assert native_layer._k_scale_float == 1.0
     assert native_layer._v_scale_float == 1.0

tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py

@@ -10,12 +10,13 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache
 from vllm.v1.worker.gpu.block_table import BlockTables
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
+from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionAdapter,
     DiffusionPagedAttentionLayerAdapter,
@@ -23,6 +24,7 @@ from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionRowBinding,
 )
 from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig
+from tests.helpers.kv_layout import build_kv_cache_tensor
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 
@@ -89,10 +91,7 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
     kv_cache_config = KVCacheConfig(
         num_blocks=4,
         kv_cache_tensors=[
-            KVCacheTensor(
-                size=canonical_spec.page_size_bytes * 4,
-                shared_by=[_LAYER_NAME],
-            )
+            build_kv_cache_tensor(canonical_spec, 4, [_LAYER_NAME])
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -101,6 +100,9 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
             )
         ],
     )
+    # 0.29 reads the resolved physical layout when allocating; the diffusion
+    # control plane pins it, so this standalone harness must too.
+    resolve_diffusion_kv_cache_layout(vllm_config)
     with set_current_vllm_config(vllm_config):
         attn_groups, _, kernel_block_sizes = init_attn_backend(
             kv_cache_config,
@@ -108,13 +110,13 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
             device,
         )
         runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
+        # vLLM 0.29 dropped attn_groups and cache_dtype here: the layout now
+        # comes from the resolved CacheConfig.kv_cache_layout inside the call.
         init_kv_cache(
             runner_kv_caches,
             vllm_config.compilation_config.static_forward_context,
             kv_cache_config,
-            attn_groups,
             device,
-            vllm_config.cache_config.cache_dtype,
             kernel_block_sizes,
             vllm_config,
         )

tests/diffusion/diffusion_kv/test_request.py

@@ -10,11 +10,11 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
-    KVCacheTensor,
 )
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
+from tests.helpers.kv_layout import build_kv_cache_tensor
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -119,7 +119,7 @@ def _manager(*, enable_caching: bool) -> KVCacheManager:
     group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)
     config = KVCacheConfig(
         num_blocks=8,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["layer0"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["layer0"])],
         kv_cache_groups=[group],
     )
     return KVCacheManager(

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -8,13 +8,14 @@ from unittest.mock import Mock
 import pytest
 import torch
 from torch import nn
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
 import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+from tests.helpers.kv_layout import build_kv_cache_tensor, layout_for_backend
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -79,7 +80,9 @@ def test_runner_discovers_native_spec_from_loaded_attention() -> None:
     assert spec.num_kv_heads == 2
     assert spec.head_size == 8
     assert spec.dtype is torch.bfloat16
-    assert spec.indexes_kv_by_block_stride is True
+    # vLLM 0.29 moved the block-stride decision off the spec onto the resolved
+    # physical layout; the backend's declaration must still select one.
+    assert layout_for_backend(_Backend).is_block_outermost is True
     assert spec.non_causal is True
 
 
@@ -112,7 +115,7 @@ def test_runner_retains_matching_rank_local_config() -> None:
     spec = runner.get_kv_cache_spec()["image_attention"]
     config = KVCacheConfig(
         num_blocks=8,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["image_attention"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)],
     )
 
@@ -126,7 +129,7 @@ def test_runner_rejects_rank_local_config_for_different_layers() -> None:
     spec = runner.get_kv_cache_spec()["image_attention"]
     config = KVCacheConfig(
         num_blocks=8,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["other_attention"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)],
     )
 

tests/diffusion/test_diffusion_scheduler.py

@@ -10,7 +10,7 @@ import pytest
 import torch
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
 from pytest_mock import MockerFixture
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
@@ -27,6 +27,7 @@ from vllm_omni.diffusion.sched import (
 from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
 from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from tests.helpers.kv_layout import build_kv_cache_tensor
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -104,7 +105,7 @@ def _initialize_paged_scheduler(
     )
     config = KVCacheConfig(
         num_blocks=num_blocks,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
     )
     scheduler.initialize(

tests/helpers/kv_layout.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared access to the production ``KVCacheTensor`` layout derivation.
+
+Tests deliberately go through the same code production uses
+(``vllm_omni.diffusion.diffusion_kv.layout``) rather than recomputing strides,
+so a regression in that derivation fails the suite instead of being masked by a
+parallel test-only copy.
+"""
+
+from types import SimpleNamespace
+
+from vllm.config import CacheConfig
+from vllm.v1.kv_cache_layout import KVCacheLayout
+
+from vllm_omni.diffusion.diffusion_kv.layout import (
+    build_kv_cache_tensor,
+    resolve_diffusion_kv_cache_layout,
+)
+
+__all__ = ["build_kv_cache_tensor", "layout_for_backend"]
+
+
+def layout_for_backend(attn_backend) -> KVCacheLayout:
+    """The physical layout ``attn_backend``'s block-stride declaration selects.
+
+    Before 0.29 a backend's ``indexes_kv_by_block_stride()`` was copied onto
+    ``AttentionSpec``; asserting on the resolved layout is the surviving way to
+    check that declaration still reaches the cache.
+    """
+
+    vllm_config = SimpleNamespace(cache_config=CacheConfig())
+    return resolve_diffusion_kv_cache_layout(
+        vllm_config,
+        indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(),
+    )

vllm_omni/diffusion/attention/backends/abstract.py

@@ -121,6 +121,13 @@ class AttentionBackend(ABC):
         when it unifies cache layouts across layers. Dense diffusion backends
         conservatively keep the default ``False``; a paged backend should
         override this only when its kernel actually follows that layout.
+
+        Since vLLM 0.29 this no longer rides on ``AttentionSpec``.  It selects
+        the block-outermost physical layout pinned on
+        ``CacheConfig.kv_cache_layout`` -- see
+        ``vllm_omni.diffusion.diffusion_kv.layout`` -- whose
+        ``is_block_outermost`` is what upstream reads when it sizes
+        ``KVCacheTensor`` regions.
         """
 
         return False

vllm_omni/diffusion/attention/layer.py

@@ -24,6 +24,7 @@ from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.layout import assert_backend_layout_supported
 from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
 from vllm_omni.diffusion.forward_context import (
     get_forward_context,
@@ -261,15 +262,19 @@ class Attention(nn.Module):
             return None
         dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype
         # Keep backend layout discovery under the same config context used by
-        # upstream vLLM's attention-spec collector.
+        # upstream vLLM's attention-spec collector.  vLLM 0.29 moved the
+        # block-stride decision off the spec and onto the single physical
+        # ``CacheConfig.kv_cache_layout`` resolved before the cache is built,
+        # so the backend's preference is enforced there (see
+        # ``vllm_omni.diffusion.diffusion_kv.initialization``) rather than
+        # carried per layer.
         with set_current_vllm_config(vllm_config):
-            indexes_kv_by_block_stride = self.attn_backend.indexes_kv_by_block_stride()
+            assert_backend_layout_supported(vllm_config, self.attn_backend)
         return FullAttentionSpec(
             block_size=vllm_config.cache_config.block_size,
             num_kv_heads=self.num_kv_heads,
             head_size=self.head_size,
             dtype=dtype,
-            indexes_kv_by_block_stride=indexes_kv_by_block_stride,
             non_causal=not self.causal,
         )
 

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

vllm_omni/experimental/ar_diffusion/kv_cache/manager.py

@@ -22,10 +22,10 @@ from vllm.v1.kv_cache_interface import (
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheSpec,
-    KVCacheTensor,
 )
 from vllm.v1.request import RequestStatus
 
+from vllm_omni.diffusion.diffusion_kv.layout import build_kv_cache_tensor
 from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
 from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
@@ -134,7 +134,7 @@ def build_kv_manager(
     """
     layer_names = list(layer_names)
     group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)
-    tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)]
+    tensors = [build_kv_cache_tensor(spec, num_blocks, layer_names)]
     config = KVCacheConfig(
         num_blocks=num_blocks,
         kv_cache_tensors=tensors,