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(