diff --git a/tests/diffusion/diffusion_kv/test_layout.py b/tests/diffusion/diffusion_kv/test_layout.py index 713849e95..e82772ea1 100644 --- a/tests/diffusion/diffusion_kv/test_layout.py +++ b/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) diff --git a/vllm_omni/diffusion/diffusion_kv/initialization.py b/vllm_omni/diffusion/diffusion_kv/initialization.py index d1a79d934..bac0a4d96 100644 --- a/vllm_omni/diffusion/diffusion_kv/initialization.py +++ b/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) diff --git a/vllm_omni/diffusion/diffusion_kv/layout.py b/vllm_omni/diffusion/diffusion_kv/layout.py index 960df395d..0af4f5d3f 100644 --- a/vllm_omni/diffusion/diffusion_kv/layout.py +++ b/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.