[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/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"):