[Bugfix] Fix native Mooncake cross-stage KV layout mismatch (#8003)

Signed-off-by: Acerak01-fy <wfy2003324@163.com>

tests/diffusion/diffusion_kv/test_layout.py

@@ -17,6 +17,7 @@ from vllm_omni.diffusion.diffusion_kv.layout import (
     adopt_kv_cache_layout,
     assert_backend_layout_supported,
     build_kv_cache_tensor,
+    get_connector_required_kv_cache_layout,
     resolve_diffusion_kv_cache_layout,
 )
 
@@ -135,6 +136,51 @@ def test_block_stride_requirement_propagates_control_plane_to_worker() -> None:
         assert_backend_layout_supported(worker, _DenseBackend)
 
 
+def test_explicit_connector_layout_is_distinct_from_block_stride_requirement() -> None:
+    config = _config()
+
+    layout = resolve_diffusion_kv_cache_layout(config, required_layout=KVCacheLayout.LBHNC)
+
+    assert layout is KVCacheLayout.LBHNC
+    assert config.cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.LBHNC
+
+
+def test_connector_declared_layout_is_resolved(monkeypatch: pytest.MonkeyPatch) -> None:
+    class _Connector:
+        @classmethod
+        def get_required_kvcache_layout(cls, config):
+            del config
+            return "LBHNC"
+
+    from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+
+    monkeypatch.setattr(KVConnectorFactory, "get_connector_class", lambda config: _Connector)
+    config = _config()
+    config.kv_transfer_config = object()
+
+    assert get_connector_required_kv_cache_layout(config) is KVCacheLayout.LBHNC
+
+
+def test_connector_layout_rejects_same_block_order_but_different_head_order() -> None:
+    config = _config()
+    config.cache_config.kv_cache_layout = KVCacheLayout.LBNHC.name
+    with pytest.raises(ValueError, match="contradicts"):
+        resolve_diffusion_kv_cache_layout(config, required_layout=KVCacheLayout.LBHNC)
+
+
+def test_connector_layout_cannot_override_backend_requirement() -> None:
+    with pytest.raises(ValueError, match="contradicts the attention"):
+        resolve_diffusion_kv_cache_layout(
+            _config(), indexes_kv_by_block_stride=True, required_layout=KVCacheLayout.LBHNC
+        )
+
+
+def test_connector_layout_without_cache_config() -> None:
+    config = _config()
+    config.cache_config = None
+    assert resolve_diffusion_kv_cache_layout(config, required_layout="LBHNC") is KVCacheLayout.LBHNC
+
+
 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)
 

vllm_omni/diffusion/diffusion_kv/initialization.py

@@ -16,7 +16,10 @@ 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.diffusion_kv.layout import (
+    get_connector_required_kv_cache_layout,
+    resolve_diffusion_kv_cache_layout,
+)
 from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
 from vllm_omni.platforms import current_omni_platform
 
@@ -49,6 +52,7 @@ def build_native_kv_cache_configs(
     layout = resolve_diffusion_kv_cache_layout(
         vllm_config,
         indexes_kv_by_block_stride=indexes_kv_by_block_stride,
+        required_layout=get_connector_required_kv_cache_layout(vllm_config),
     )
     single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config)
     worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory)

vllm_omni/diffusion/diffusion_kv/layout.py

@@ -39,6 +39,7 @@ def resolve_diffusion_kv_cache_layout(
     vllm_config: VllmConfig,
     *,
     indexes_kv_by_block_stride: bool = False,
+    required_layout: KVCacheLayout | str | None = None,
 ) -> KVCacheLayout:
     """Pin the physical KV layout for the Diffusion cache and return it.
 
@@ -48,31 +49,67 @@ def resolve_diffusion_kv_cache_layout(
     from the wrong offsets instead of raising.
     """
 
+    if required_layout is not None:
+        required_layout = KVCacheLayout[required_layout] if isinstance(required_layout, str) else required_layout
+        if not isinstance(required_layout, KVCacheLayout):
+            raise TypeError(f"required_layout must be a KVCacheLayout or name, got {type(required_layout)!r}")
     required_block_outermost = bool(indexes_kv_by_block_stride)
+    if required_layout is not None and required_layout.is_block_outermost != required_block_outermost:
+        raise ValueError(
+            f"Connector KV cache layout {required_layout.name} contradicts the attention "
+            f"backend, which requires is_block_outermost={required_block_outermost}."
+        )
     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
+        return required_layout or (
+            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:
+        if (required_layout is not None and layout is not required_layout) or (
+            required_layout is None and 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}."
+                "backend/connector requirement: "
+                f"layout {required_layout.name if required_layout is not None else required_block_outermost}."
             )
         return layout
 
-    layout = BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT if required_block_outermost else DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+    layout = required_layout or (
+        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 get_connector_required_kv_cache_layout(vllm_config: VllmConfig) -> KVCacheLayout | None:
+    """Resolve the physical layout declared by the configured native connector."""
+
+    transfer_config = getattr(vllm_config, "kv_transfer_config", None)
+    if transfer_config is None:
+        return None
+    from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+
+    connector_cls = KVConnectorFactory.get_connector_class(transfer_config)
+    getter = getattr(connector_cls, "get_required_kvcache_layout", None)
+    if getter is None:
+        return None
+    layout = getter(vllm_config)
+    if layout is None:
+        return None
+    try:
+        return layout if isinstance(layout, KVCacheLayout) else KVCacheLayout[layout]
+    except (KeyError, TypeError) as exc:
+        raise ValueError(f"Connector {connector_cls.__name__} returned unsupported KV layout {layout!r}") from exc
+
+
 def adopt_kv_cache_layout(vllm_config: VllmConfig, kv_cache_config) -> KVCacheLayout:
     """Adopt the control plane's resolved layout inside a worker.