[Diffusion] Implement paged AR↔DiT KV connector ("v1", reusing vllm's native mooncake) (#6310)

Signed-off-by: asukaqaq-s <1311722138@qq.com>
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/config/test_omni_config.py

@@ -1269,6 +1269,10 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
                 "    fa_deterministic: true",
                 "    diffusion_kv_mode: paged_scheduler",
                 "    diffusion_kv_max_rows_per_request: 2",
+                "    kv_transfer_config:",
+                "      kv_connector: MooncakeConnector",
+                "      kv_role: kv_consumer",
+                "      engine_id: dit-engine-1",
             ]
         )
     )
@@ -1283,6 +1287,14 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
     assert stage.diffusion_config.fa_deterministic is True
     assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
     assert stage.diffusion_config.diffusion_kv_max_rows_per_request == 2
+    assert stage.diffusion_config.kv_transfer_config.engine_id == "dit-engine-1"
+
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
+
+    engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
+    od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
+    assert od_config.kv_transfer_config.engine_id == "dit-engine-1"
 
 
 def test_from_pipeline_config_forwards_fastvideo_vsa_topk(tmp_path, monkeypatch):

tests/diffusion/diffusion_kv/test_initialization.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from types import SimpleNamespace
 from unittest.mock import Mock

tests/diffusion/diffusion_kv/test_kv_connector.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from unittest import mock
+
+import pytest
+import torch
+from vllm.config import KVTransferConfig
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_kv.kv_connector import (
+    parse_kv_transfer_config,
+    shutdown_kv_connector,
+)
+from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler
+from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+KV_TRANSFER_CONFIG = {
+    "kv_connector": "MooncakeConnector",
+    "kv_role": "kv_consumer",
+    "engine_id": "dit-engine-1",
+    "kv_connector_extra_config": {"mooncake_protocol": "tcp"},
+}
+
+
+@pytest.fixture(autouse=True)
+def _fixed_master_port(monkeypatch) -> None:
+    monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
+
+
+def test_config_roundtrip_to_kv_vllm_config() -> None:
+    od_config = OmniDiffusionConfig.from_kwargs(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=1,
+        max_model_len=64,
+        kv_transfer_config=dict(KV_TRANSFER_CONFIG),
+    )
+
+    assert isinstance(od_config.kv_transfer_config, KVTransferConfig)
+    assert od_config.kv_transfer_config.engine_id == "dit-engine-1"
+    vllm_config = create_diffusion_vllm_config(torch.device("cpu"), od_config)
+    assert vllm_config.kv_transfer_config is od_config.kv_transfer_config
+
+
+def test_parse_requires_explicit_engine_id() -> None:
+    payload = dict(KV_TRANSFER_CONFIG)
+    payload.pop("engine_id")
+    with pytest.raises(ValueError, match="non-empty engine_id"):
+        parse_kv_transfer_config(payload)
+
+
+def test_diffusion_projection_rejects_missing_engine_id_before_materialize() -> None:
+    """Structured resolve must not let upstream auto-mint engine_id first."""
+    from vllm_omni.config.omni_config import _DiffusionConfigProjection
+
+    payload = dict(KV_TRANSFER_CONFIG)
+    payload.pop("engine_id")
+    with pytest.raises(ValueError, match="non-empty engine_id"):
+        _DiffusionConfigProjection.from_kwargs(kv_transfer_config=payload)
+
+
+def test_diffusion_projection_preserves_explicit_engine_id() -> None:
+    from vllm_omni.config.omni_config import _DiffusionConfigProjection
+
+    projection = _DiffusionConfigProjection.from_kwargs(kv_transfer_config=dict(KV_TRANSFER_CONFIG))
+    assert isinstance(projection.kv_transfer_config, KVTransferConfig)
+    assert projection.kv_transfer_config.engine_id == "dit-engine-1"
+
+    od_config = OmniDiffusionConfig.from_kwargs(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=1,
+        kv_transfer_config=projection.kv_transfer_config,
+    )
+    assert od_config.kv_transfer_config is not None
+    assert od_config.kv_transfer_config.engine_id == "dit-engine-1"
+
+
+class _ConcreteScheduler(BaseScheduler):
+    def update_from_output(self, sched_output, output) -> set[str]:
+        del sched_output, output
+        return set()
+
+
+def test_scheduler_assembles_kv_stub_and_shuts_it_down() -> None:
+    scheduler = _ConcreteScheduler()
+    fake_connector = mock.Mock()
+    od_config = OmniDiffusionConfig.from_kwargs(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=1,
+        max_model_len=64,
+        kv_transfer_config=dict(KV_TRANSFER_CONFIG),
+    )
+    kv_vllm_config = mock.Mock()
+    kv_vllm_config.model_config.max_model_len = 64
+    kv_vllm_config.max_in_flight_tokens = 64
+
+    with (
+        mock.patch(
+            "vllm_omni.diffusion.diffusion_kv.kv_connector.KVConnectorFactory.create_connector",
+            return_value=fake_connector,
+        ),
+        mock.patch("vllm_omni.diffusion.sched.base_scheduler.DiffusionKVCacheManager"),
+    ):
+        scheduler.initialize(
+            od_config,
+            kv_cache_config=mock.sentinel.kv_cache_config,
+            scheduler_block_size=16,
+            hash_block_size=16,
+            kv_vllm_config=kv_vllm_config,
+        )
+
+    assert scheduler.kv_connector is fake_connector
+    with mock.patch("vllm_omni.diffusion.diffusion_kv.kv_connector.shutdown_kv_connector") as shutdown:
+        scheduler.close()
+    shutdown.assert_called_once_with(scheduler_connector=fake_connector)
+
+
+def test_native_config_requires_paged_scheduler() -> None:
+    with pytest.raises(ValueError, match="requires diffusion_kv_mode='paged_scheduler'"):
+        OmniDiffusionConfig.from_kwargs(kv_transfer_config=dict(KV_TRANSFER_CONFIG))
+
+
+def test_kv_and_legacy_transfer_configs_are_exclusive() -> None:
+    with pytest.raises(ValueError, match="exactly one KV transfer path"):
+        OmniDiffusionConfig.from_kwargs(
+            kv_transfer_config=dict(KV_TRANSFER_CONFIG),
+            omni_kv_config={"need_recv_cache": True},
+        )
+
+
+def test_shutdown_kv_connector_is_idempotent() -> None:
+    scheduler_connector = mock.Mock()
+    with mock.patch("vllm_omni.diffusion.diffusion_kv.kv_connector.ensure_kv_transfer_shutdown") as ensure_shutdown:
+        shutdown_kv_connector(scheduler_connector=scheduler_connector)
+        shutdown_kv_connector()
+
+    assert ensure_shutdown.call_count == 2
+    scheduler_connector.shutdown.assert_called_once()

tests/diffusion/diffusion_kv/test_metadata.py

@@ -8,10 +8,83 @@ from vllm_omni.diffusion.diffusion_kv.metadata import (
     DiffusionKVMetadata,
     DiffusionKVSequenceMetadata,
 )
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler
+from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
 
 
+def test_native_request_metadata_round_trips_and_stays_optional() -> None:
+    params = {
+        "transfer_id": "xfer-req-0",
+        "do_remote_decode": True,
+        "do_remote_prefill": False,
+        "remote_engine_id": "dit-engine-0",
+    }
+    kv_request = DiffusionKVRequest(
+        "req-0/diffusion-kv/0",
+        sequence_id=0,
+        prefix_len=4,
+        target_len=2,
+        seq_len=6,
+        kv_transfer_params=params,
+        prompt_token_ids=[0] * 6,
+    )
+    public_request = OmniDiffusionRequest(
+        prompt="prompt",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_id="req-0",
+        kv_transfer_params=params,
+    )
+    scheduler_output = DiffusionSchedulerOutput(
+        step_id=0,
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        finished_req_ids=set(),
+        num_running_reqs=0,
+        num_waiting_reqs=0,
+    )
+
+    assert kv_request.kv_transfer_params is params
+    assert kv_request.prompt_token_ids == [0] * 6
+    assert public_request.kv_transfer_params is params
+    assert scheduler_output.kv_connector_metadata is None
+
+
+class _MetadataScheduler(BaseScheduler):
+    def update_from_output(self, sched_output, output) -> set[str]:
+        del sched_output, output
+        return set()
+
+
+def test_scheduler_copies_opaque_params_to_sequence_requests() -> None:
+    params = {"transfer_id": "xfer-req-0", "do_remote_prefill": True}
+    sequence = DiffusionKVRequest(
+        "req-0/diffusion-kv/0",
+        sequence_id=0,
+        prefix_len=4,
+        target_len=2,
+        seq_len=6,
+    )
+    request = OmniDiffusionRequest(
+        prompt="prompt",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_id="req-0",
+        diffusion_kv_requests=(sequence,),
+        kv_transfer_params=params,
+    )
+    scheduler = _MetadataScheduler()
+    scheduler._diffusion_kv_manager = object()
+
+    state = scheduler._make_request_state("req-0", request)
+
+    assert state.diffusion_kv_requests[0].kv_transfer_params == params
+    assert state.diffusion_kv_requests[0].kv_transfer_params is not params
+
+
 def test_diffusion_kv_metadata_uses_native_cache_group_block_ids() -> None:
     context = DiffusionKVContextMetadata(
         context_id="text",

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -134,21 +134,28 @@ def test_runner_rejects_rank_local_config_for_different_layers() -> None:
         runner.set_kv_cache_config(config)
 
 
-def test_worker_selects_its_rank_local_config() -> None:
+def test_worker_selects_its_rank_local_config(monkeypatch) -> None:
     worker = object.__new__(DiffusionWorker)
     worker.rank = 1
     worker.od_config = SimpleNamespace(num_gpus=2)
     worker.vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(max_model_len=None),
         cache_config=SimpleNamespace(num_gpu_blocks=None),
+        kv_transfer_config=SimpleNamespace(kv_connector="MooncakeConnector", engine_id="dit-engine-1"),
     )
     worker._maybe_get_memory_pool_context = lambda _tag: nullcontext()
     worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config))
     configs = [SimpleNamespace(num_blocks=4), SimpleNamespace(num_blocks=8)]
+    ensure_initialized = Mock()
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_kv.kv_connector.ensure_kv_transfer_initialized",
+        ensure_initialized,
+    )
 
     worker.set_kv_cache_configs(configs, 64)
 
     assert worker.installed is configs[1]
+    ensure_initialized.assert_called_once_with(worker.vllm_config, configs[1])
     assert worker.vllm_config.model_config.max_model_len == 64
     assert worker.vllm_config.cache_config.num_gpu_blocks == 8
     with pytest.raises(ValueError, match="rank count mismatch"):

vllm_omni/config/omni_config.py

@@ -18,10 +18,11 @@ from inspect import Parameter, signature
 from pathlib import Path
 from typing import Any, Literal, TypeAlias, TypedDict, cast
 
-from pydantic import ConfigDict, Field, model_validator
+from pydantic import ConfigDict, Field, field_validator, model_validator
 from typing_extensions import Self
 from vllm.config import CacheConfig as VllmCacheConfig
 from vllm.config import CompilationConfig as VllmCompilationConfig
+from vllm.config import KVTransferConfig
 from vllm.config import LoadConfig as VllmLoadConfig
 from vllm.config import ParallelConfig as VllmParallelConfig
 from vllm.config import ProfilerConfig as VllmProfilerConfig
@@ -775,11 +776,19 @@ class _DiffusionConfigProjection:
     worker_extension_cls: str | None = None
     custom_pipeline_args: dict[str, Any] | None = None
     additional_config: dict[str, Any] = field(default_factory=dict)
+    kv_transfer_config: KVTransferConfig | None = None
     enable_stage_verification: bool = True
     prompt_file_path: str | None = None
     quantization_config: _QuantizationConfigType = None
     extras: dict[str, Any] = field(default_factory=dict)
 
+    @field_validator("kv_transfer_config", mode="before")
+    @classmethod
+    def _normalize_kv_transfer_config(cls, value: Any) -> Any:
+        from vllm_omni.diffusion.diffusion_kv.kv_connector import parse_kv_transfer_config
+
+        return parse_kv_transfer_config(value)
+
     @classmethod
     def from_kwargs(cls, **kwargs: Any) -> _DiffusionConfigProjection:
         from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs

vllm_omni/diffusion/data.py

@@ -34,7 +34,7 @@ from vllm_omni.errors import client_error_metadata
 from vllm_omni.quantization import build_quant_config
 
 if TYPE_CHECKING:
-    from vllm.config import ProfilerConfig
+    from vllm.config import KVTransferConfig, ProfilerConfig
 
 # Import after TYPE_CHECKING to avoid circular imports at runtime
 # The actual import is deferred to __post_init__ to avoid import order issues
@@ -935,6 +935,9 @@ class OmniDiffusionConfig:
 
     # Omni configuration (injected from stage config)
     omni_kv_config: dict[str, Any] = field(default_factory=dict)
+    # Native vLLM KV connector configuration. PR0 only assembles the connector;
+    # the native page data path is owned by the follow-up landing PR.
+    kv_transfer_config: "KVTransferConfig | None" = None
     additional_config: dict[str, Any] = field(default_factory=dict)
 
     profiler_config: "ProfilerConfig | dict[str, Any] | None" = None
@@ -1105,6 +1108,19 @@ class OmniDiffusionConfig:
                 "paged_scheduler Diffusion KV does not support imported AR KV; "
                 "disable need_recv_cache until connector-aware import is implemented"
             )
+
+        if self.kv_transfer_config is not None:
+            from vllm_omni.diffusion.diffusion_kv.kv_connector import parse_kv_transfer_config
+
+            if any(self.omni_kv_config.get(name, False) for name in ("need_send_cache", "need_recv_cache")):
+                raise ValueError(
+                    "native kv_transfer_config cannot be combined with legacy omni_kv_config transfer; "
+                    "configure exactly one KV transfer path"
+                )
+            if self.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
+                raise ValueError("native kv_transfer_config requires diffusion_kv_mode='paged_scheduler'")
+            self.kv_transfer_config = parse_kv_transfer_config(self.kv_transfer_config)
+
         self.master_port = self._resolve_master_port()
         self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0)
         if not math.isfinite(self.request_batch_max_wait_ms) or self.request_batch_max_wait_ms < 0:

vllm_omni/diffusion/diffusion_kv/kv_connector.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Diffusion-side assembly facade built on the upstream Mooncake KV connector.
+
+The configured ``kv_connector`` (currently ``MooncakeConnector``) is created
+through vLLM's native ``KVConnectorFactory``. PR0 only validates configuration
+and creates the Scheduler/Worker connector objects; it does not subclass or
+reimplement Mooncake's transfer runtime. Real page sizing, cache registration,
+admission, and remote loading belong to the follow-up landing PR.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any
+
+import torch
+from vllm.config import KVTransferConfig, VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+from vllm.distributed.kv_transfer.kv_transfer_state import ensure_kv_transfer_initialized, ensure_kv_transfer_shutdown
+from vllm.logger import init_logger
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
+if TYPE_CHECKING:
+    from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1
+
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+logger = init_logger(__name__)
+
+
+def mint_transfer_id(request_id: str) -> str:
+    """Return the stable ticket shared by one AR -> Diffusion handoff."""
+
+    return f"xfer-{request_id}"
+
+
+def build_source_kv_transfer_params(
+    *,
+    transfer_id: str,
+    remote_engine_id: str | None,
+    remote_bootstrap_addr: str | None,
+) -> dict[str, Any]:
+    """Build the opaque producer-side metadata bag for an AR request."""
+
+    params: dict[str, Any] = {
+        "transfer_id": transfer_id,
+        "do_remote_decode": True,
+        "do_remote_prefill": False,
+    }
+    if remote_engine_id:
+        params["remote_engine_id"] = remote_engine_id
+    if remote_bootstrap_addr:
+        params["remote_bootstrap_addr"] = remote_bootstrap_addr
+    return params
+
+
+def build_target_kv_transfer_params(
+    *,
+    source_params: Mapping[str, Any],
+    remote_engine_id: str | None,
+    remote_bootstrap_addr: str | None,
+) -> dict[str, Any]:
+    """Build the consumer-side metadata bag without interpreting pages."""
+
+    params = dict(source_params)
+    params["do_remote_prefill"] = True
+    params["do_remote_decode"] = False
+    if remote_engine_id:
+        params["remote_engine_id"] = remote_engine_id
+    if remote_bootstrap_addr:
+        params["remote_bootstrap_addr"] = remote_bootstrap_addr
+    return params
+
+
+def bootstrap_addr_from_kv_transfer_config(kv_transfer_config: KVTransferConfig | None) -> str | None:
+    """Read an optional connector bootstrap endpoint from native config."""
+
+    if kv_transfer_config is None:
+        return None
+    extra_config = kv_transfer_config.kv_connector_extra_config or {}
+    for key in ("bootstrap_addr", "prefill_bootstrap_addr", "mooncake_master"):
+        value = extra_config.get(key)
+        if isinstance(value, str) and value.strip():
+            return value
+    return None
+
+
+def parse_kv_transfer_config(value: object | None) -> KVTransferConfig | None:
+    """Normalize a YAML mapping or an existing native config."""
+
+    if value is None:
+        return None
+    if isinstance(value, KVTransferConfig):
+        _validate_kv_transfer_config(value)
+        return value
+    if isinstance(value, Mapping):
+        payload = dict(value)
+        if not payload:
+            return None
+        engine_id = payload.get("engine_id")
+        if not isinstance(engine_id, str) or not engine_id.strip():
+            raise ValueError("Diffusion native kv_transfer_config requires a non-empty engine_id")
+        config = KVTransferConfig(**payload)
+        _validate_kv_transfer_config(config)
+        return config
+    raise TypeError(f"kv_transfer_config must be a mapping or KVTransferConfig, got {type(value)!r}")
+
+
+def create_scheduler_kv_connector(
+    od_config: OmniDiffusionConfig,
+) -> KVConnectorBase_V1 | None:
+    """Create a Scheduler-role connector when native config is present."""
+
+    kv_transfer_config = getattr(od_config, "kv_transfer_config", None)
+    if kv_transfer_config is None:
+        return None
+    if not isinstance(kv_transfer_config, KVTransferConfig):
+        raise TypeError(
+            f"Diffusion native kv_transfer_config must be KVTransferConfig, got {type(kv_transfer_config)!r}"
+        )
+
+    vllm_config = _build_native_vllm_config(od_config)
+    connector = KVConnectorFactory.create_connector(
+        config=vllm_config,
+        role=KVConnectorRole.SCHEDULER,
+        kv_cache_config=_empty_kv_cache_config(),
+    )
+    logger.info(
+        "Created KV connector stub (SCHEDULER role): connector=%s engine_id=%s",
+        kv_transfer_config.kv_connector,
+        kv_transfer_config.engine_id,
+    )
+    return connector
+
+
+def init_worker_kv_connector(vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None:
+    """Initialize the Worker-role connector with its rank-local cache plan."""
+
+    if vllm_config.kv_transfer_config is None:
+        return
+    ensure_kv_transfer_initialized(vllm_config, kv_cache_config)
+    logger.info(
+        "Initialized KV connector stub (WORKER role): connector=%s engine_id=%s",
+        vllm_config.kv_transfer_config.kv_connector,
+        vllm_config.kv_transfer_config.engine_id,
+    )
+
+
+def shutdown_kv_connector(*, scheduler_connector: KVConnectorBase_V1 | None = None) -> None:
+    """Shutdown Worker and Scheduler connector objects idempotently."""
+
+    if scheduler_connector is not None:
+        scheduler_connector.shutdown()
+    ensure_kv_transfer_shutdown()
+
+
+def _empty_kv_cache_config() -> KVCacheConfig:
+    """Empty cache config used only for PR0 connector assembly."""
+
+    return KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[])
+
+
+def _build_native_vllm_config(od_config: OmniDiffusionConfig) -> VllmConfig:
+    from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
+
+    kv_transfer_config = getattr(od_config, "kv_transfer_config", None)
+    assert isinstance(kv_transfer_config, KVTransferConfig)
+    vllm_config = create_diffusion_vllm_config(torch.device("cpu"), od_config)
+    vllm_config.kv_transfer_config = kv_transfer_config
+    return vllm_config
+
+
+def _validate_kv_transfer_config(config: KVTransferConfig) -> None:
+    engine_id = config.engine_id
+    if not isinstance(engine_id, str) or not engine_id.strip():
+        raise ValueError("Diffusion native kv_transfer_config requires a non-empty engine_id")
+    if config.kv_connector is not None and config.kv_role is None:
+        raise ValueError("Diffusion native kv_transfer_config requires kv_role when kv_connector is set")
+    if config.kv_role not in (None, "kv_consumer", "kv_producer", "kv_both"):
+        raise ValueError(f"Unsupported KV connector role: {config.kv_role!r}")

vllm_omni/diffusion/diffusion_kv/request.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 from dataclasses import dataclass
+from typing import Any
 
 from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.request import RequestStatus
@@ -64,6 +65,8 @@ class DiffusionKVRequest:
         seq_len: int,
         block_hashes: Sequence[BlockHash] = (),
         kv_contexts: Sequence[DiffusionKVContext] = (),
+        kv_transfer_params: dict[str, Any] | None = None,
+        prompt_token_ids: list[int] | None = None,
     ) -> None:
         if not request_id:
             raise ValueError("request_id must be non-empty")
@@ -108,6 +111,11 @@ class DiffusionKVRequest:
         # vLLM 0.26+ uses this optional boundary when publishing cached blocks;
         # zero means that no sparse-retention boundary is active.
         self.shared_prefix_boundary = 0
+        # Opaque request-scoped parameters consumed by an upstream KVConnector.
+        self.kv_transfer_params = kv_transfer_params
+        # Mooncake derives the remote-load length from this list. PR0 carries
+        # the field without creating a placeholder or starting a load.
+        self.prompt_token_ids = prompt_token_ids
 
     @property
     def seq_len(self) -> int:

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -107,6 +107,7 @@ class InlineStageDiffusionClient(StageClientBase):
         prompt: OmniPromptType,
         sampling_params: OmniDiffusionSamplingParams,
         kv_sender_info: dict[int, dict[str, Any]] | None = None,
+        kv_transfer_params: dict[str, Any] | None = None,
     ) -> None:
         # Each request mutates its sampling state while it is normalized and
         # executed. Callers commonly reuse one params object for concurrent
@@ -124,6 +125,7 @@ class InlineStageDiffusionClient(StageClientBase):
                 prompt,
                 sampling_params,
                 kv_sender_info,
+                kv_transfer_params,
             )
         )
         self._tasks[request_id] = task
@@ -134,6 +136,7 @@ class InlineStageDiffusionClient(StageClientBase):
         prompt: Any,
         sampling_params: OmniDiffusionSamplingParams,
         kv_sender_info: dict[str, Any] | None = None,
+        kv_transfer_params: dict[str, Any] | None = None,
     ) -> None:
         try:
             request = OmniDiffusionRequest(
@@ -141,6 +144,7 @@ class InlineStageDiffusionClient(StageClientBase):
                 sampling_params=sampling_params,
                 request_id=request_id,
                 kv_sender_info=kv_sender_info,
+                kv_transfer_params=kv_transfer_params,
             )
 
             if self.od_config.streaming_output:

vllm_omni/diffusion/request.py

@@ -67,6 +67,9 @@ class OmniDiffusionRequest:
     # This is populated by a pipeline preprocessor before the request reaches
     # the scheduler; ``None`` keeps the default behavior for other pipelines.
     batch_compatibility_key: tuple[Any, ...] | None = None
+    # Opaque native vLLM connector parameters. The Orchestrator and Scheduler
+    # transport this bag without interpreting local pages or block IDs.
+    kv_transfer_params: dict[str, Any] | None = None
     # A model preprocessor may keep selected requests on the legacy full-forward
     # path even when the engine is globally configured for step execution.
     use_step_execution: bool = True

vllm_omni/diffusion/sched/base_scheduler.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -57,6 +57,7 @@ class BaseScheduler(ABC):
         self.max_num_running_reqs: int = 1
         self._prefetch_enabled: bool = False
         self._diffusion_kv_manager: DiffusionKVCacheManager | None = None
+        self._kv_connector = None
 
     def initialize(
         self,
@@ -111,8 +112,17 @@ class BaseScheduler(ABC):
             ):
                 raise ValueError("dense_legacy Scheduler received unexpected Diffusion KV cache initialization state")
             self._diffusion_kv_manager = None
+        from vllm_omni.diffusion.diffusion_kv.kv_connector import create_scheduler_kv_connector
+
+        self._kv_connector = create_scheduler_kv_connector(od_config)
         self._reset_scheduler_state()
 
+    @property
+    def kv_connector(self):
+        """Upstream vLLM Scheduler-role connector, when configured."""
+
+        return self._kv_connector
+
     def add_request(self, request: OmniDiffusionRequest) -> str:
         return self._add_request_with_request_id(request.request_id, request)
 
@@ -288,6 +298,10 @@ class BaseScheduler(ABC):
         if self._diffusion_kv_manager is not None:
             self._diffusion_kv_manager.close()
             self._diffusion_kv_manager = None
+        from vllm_omni.diffusion.diffusion_kv.kv_connector import shutdown_kv_connector
+
+        shutdown_kv_connector(scheduler_connector=self._kv_connector)
+        self._kv_connector = None
         self._request_states.clear()
         self._waiting.clear()
         self._running.clear()
@@ -377,6 +391,14 @@ class BaseScheduler(ABC):
         elif kv_requests:
             raise ValueError("dense_legacy request unexpectedly contains Scheduler Diffusion KV requests")
 
+        # Keep the request-scoped native connector bag on each Scheduler-owned
+        # sequence without interpreting its contents or creating page state.
+        kv_transfer_params = getattr(request, "kv_transfer_params", None)
+        if kv_transfer_params is not None:
+            for kv_request in kv_requests:
+                if kv_request.kv_transfer_params is None:
+                    kv_request.kv_transfer_params = dict(kv_transfer_params)
+
         # DiffusionKVRequest objects are mutable Scheduler/native-KVCacheManager
         # state and must never ride the normal request payload to a Worker.
         request.diffusion_kv_requests = None

vllm_omni/diffusion/sched/interface.py

@@ -12,6 +12,8 @@ from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 if TYPE_CHECKING:
+    from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
+
     from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 
 
@@ -251,6 +253,9 @@ class DiffusionSchedulerOutput:
     num_waiting_reqs: int
     # next request to background-prefetch KV
     kv_prefetch_job: KVPrefetchJob | None = None
+    # Opaque metadata emitted by a future Scheduler-role connector. PR0 keeps
+    # the input port but does not build or consume it.
+    kv_connector_metadata: KVConnectorMetadata | None = None
 
     @cached_property
     def scheduled_request_ids(self) -> list[str]:

vllm_omni/diffusion/stage_diffusion_client.py

@@ -339,6 +339,7 @@ class StageDiffusionClient(StageClientBase):
         prompt: OmniPromptType,
         sampling_params: OmniDiffusionSamplingParams,
         kv_sender_info: dict[int, dict[str, Any]] | None = None,
+        kv_transfer_params: dict[str, Any] | None = None,
     ) -> None:
         if self._engine_dead:
             raise EngineDeadError()
@@ -356,6 +357,7 @@ class StageDiffusionClient(StageClientBase):
                     "prompt": prompt,
                     "sampling_params": self._sampling_params_to_dict(sampling_params),
                     "kv_sender_info": kv_sender_info,
+                    "kv_transfer_params": kv_transfer_params,
                 }
             )
         )

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -163,6 +163,7 @@ class StageDiffusionProc:
         prompt: Any,
         sampling_params_dict: dict,
         kv_sender_info: dict[str, Any] | None = None,
+        kv_transfer_params: dict[str, Any] | None = None,
     ) -> OmniRequestOutput:
         """Build a diffusion request and consume DiffusionEngine.step_streaming() to completion."""
         sampling_params = self._reconstruct_sampling_params(sampling_params_dict)
@@ -172,6 +173,7 @@ class StageDiffusionProc:
             sampling_params=sampling_params,
             request_id=request_id,
             kv_sender_info=kv_sender_info,
+            kv_transfer_params=kv_transfer_params,
         )
 
         # Non-streaming callers share the streaming engine path but only
@@ -191,6 +193,7 @@ class StageDiffusionProc:
         prompt: Any,
         sampling_params_dict: dict,
         kv_sender_info: dict[str, Any] | None = None,
+        kv_transfer_params: dict[str, Any] | None = None,
     ) -> AsyncGenerator[OmniRequestOutput, None]:
         """Process a streaming diffusion request and yield the results from DiffusionEngine.step_streaming()."""
         sampling_params = self._reconstruct_sampling_params(sampling_params_dict)
@@ -200,6 +203,7 @@ class StageDiffusionProc:
             sampling_params=sampling_params,
             request_id=request_id,
             kv_sender_info=kv_sender_info,
+            kv_transfer_params=kv_transfer_params,
         )
 
         async for results in self._engine.step_streaming(request):  # pyright: ignore[reportOptionalMemberAccess]
@@ -345,6 +349,7 @@ class StageDiffusionProc:
             prompt: Any,
             sampling_params_dict: dict,
             kv_sender_info: dict[str, Any] | None = None,
+            kv_transfer_params: dict[str, Any] | None = None,
         ) -> None:
             """Process a single diffusion request and send the response."""
             try:
@@ -354,6 +359,7 @@ class StageDiffusionProc:
                         prompt,
                         sampling_params_dict,
                         kv_sender_info=kv_sender_info,
+                        kv_transfer_params=kv_transfer_params,
                     )
                     await response_socket.send(encoder.encode({"type": "result", "output": result}))
                 else:
@@ -362,6 +368,7 @@ class StageDiffusionProc:
                         prompt,
                         sampling_params_dict,
                         kv_sender_info=kv_sender_info,
+                        kv_transfer_params=kv_transfer_params,
                     ):
                         await response_socket.send(encoder.encode({"type": "result", "output": result}))
             except DiffusionRequestAbortedError as e:
@@ -440,6 +447,7 @@ class StageDiffusionProc:
                             msg["prompt"],
                             msg["sampling_params"],
                             msg.get("kv_sender_info"),
+                            msg.get("kv_transfer_params"),
                         )
                     )
                     tasks[request_id] = task

vllm_omni/diffusion/vllm_config.py

@@ -6,10 +6,10 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any
+from typing import Any, cast
 
 import torch
-from vllm.config import CompilationConfig, DeviceConfig, VllmConfig
+from vllm.config import CompilationConfig, DeviceConfig, KVTransferConfig, VllmConfig
 from vllm.transformers_utils.config import get_hf_text_config
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -229,6 +229,10 @@ def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiff
         if max_num_batched_tokens is not None:
             vllm_config.scheduler_config.max_num_batched_tokens = int(max_num_batched_tokens)
         current_omni_platform.configure_diffusion_vllm_config(vllm_config, od_config)
+
+    kv_transfer_config = getattr(od_config, "kv_transfer_config", None)
+    if kv_transfer_config is not None:
+        vllm_config.kv_transfer_config = cast(KVTransferConfig, kv_transfer_config)
     return vllm_config
 
 

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -46,6 +46,10 @@ from vllm_omni.diffusion.data import (
     OmniWakeTask,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.kv_connector import (
+    init_worker_kv_connector,
+    shutdown_kv_connector,
+)
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -522,6 +526,7 @@ class DiffusionWorker:
         self.vllm_config.model_config.max_model_len = resolved_max_model_len
         kv_cache_config = kv_cache_configs[self.rank]
         self.vllm_config.cache_config.num_gpu_blocks = kv_cache_config.num_blocks
+        init_worker_kv_connector(self.vllm_config, kv_cache_config)
         with self._maybe_get_memory_pool_context("kv_cache"):
             self.model_runner.set_kv_cache_config(kv_cache_config)
 
@@ -978,13 +983,16 @@ class DiffusionWorker:
                         mgr.shutdown_prefetch()
         finally:
             try:
-                a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute")
-                if a2a_permute is not None:
-                    a2a_permute.clear_a2a_permute_workspaces()
-            except Exception:
-                logger.exception("Failed to release fused Ulysses symmetric-memory workspaces")
+                shutdown_kv_connector()
             finally:
-                destroy_distributed_env()
+                try:
+                    a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute")
+                    if a2a_permute is not None:
+                        a2a_permute.clear_a2a_permute_workspaces()
+                except Exception:
+                    logger.exception("Failed to release fused Ulysses symmetric-memory workspaces")
+                finally:
+                    destroy_distributed_env()
 
 
 class CustomPipelineWorkerExtension:

vllm_omni/engine/stage_client.py

@@ -104,6 +104,7 @@ class StagePoolDiffusionClient(StagePoolClient, Protocol):
         prompt: OmniPromptType,
         sampling_params: OmniDiffusionSamplingParams,
         kv_sender_info: dict[int, dict[str, Any]] | None = None,
+        kv_transfer_params: dict[str, Any] | None = None,
     ) -> None: ...
 
     def get_diffusion_output_nowait(self) -> OmniRequestOutput | None: ...