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_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)],
     )