[Feat] Add native Mooncake KV transfer from AR to DiT (#7166)

Signed-off-by: asukaqaq-s <1311722138@qq.com>
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Co-authored-by: Acerak01-fy <wfy2003324@163.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

docs/.nav.yml

@@ -161,6 +161,7 @@ nav:
         - Cache-DiT: design/feature/cache_dit.md
         - TeaCache: design/feature/teacache.md
         - Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
+        - Mooncake KV Transfer: design/feature/mooncake-kv-transfer.md
     - Infrastructure and Performance:
       - Prometheus Metrics: design/metrics.md
       - Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md

docs/design/feature/mooncake-kv-transfer.md

@@ -0,0 +1,89 @@
+# Mooncake KV Transfer
+
+## PR summary
+
+This change adds an opt-in AR-to-DiT KV handoff using vLLM's native
+`KVTransferConfig`, Scheduler-owned diffusion KV pages, and
+`MooncakeConnector`. The existing `OmniKVTransferManager` path remains the
+fallback when `kv_transfer_config` is absent.
+
+The orchestrator creates one transfer ticket per request and resolves the
+bound AR replica's Mooncake endpoint after generation. DiT allocates the
+destination pages, waits for every worker rank to finish loading, and releases
+the source and destination pages only after the connector lifecycle completes.
+Request registration failures are rolled back before Worker dispatch. Receive
+timeouts and cancellation terminate the request without recycling in-flight
+pages: the scheduler retains them until every CFG row has completed on every
+rank. Late completion is polled only while such requests remain. If a peer
+never completes, its pages remain reserved until shutdown; other requests can
+continue using the remaining capacity. Invalid pages, missing Worker ranks,
+and failures without a proven rollback still stop the replica before teardown.
+The example sets `transfer_timeout: 60.0` seconds explicitly. Cancellation
+queued during a synchronous receive is processed when that receive completes
+or times out; it does not interrupt the remote write.
+
+## Reference configuration
+
+| Stage | Replicas | Devices | Parallelism | KV cache |
+| --- | ---: | --- | --- | ---: |
+| AR | 2 | 0-3 | TP=2 per replica | 2 GiB per worker |
+| DiT | 2 | 4-7 | TP=2 per replica, EP enabled | 1 GiB per worker |
+
+The acceptance workload uses 512x512 images, two denoising steps, guidance
+scale 2.0, and `max_num_seqs=1`.
+
+For a four-GPU debug run, set both stages to one replica and assign AR TP=2 to
+devices `0,1` and DiT TP=2 to devices `2,3`.
+
+## Experiment
+
+The final comparison used one node with eight NVIDIA RTX PRO 6000 Blackwell
+Server Edition GPUs, TCP Mooncake transport, the same 50-prompt/seed T2I
+dataset, concurrency 8, and two independent 100-request runs per
+implementation. Each side completed 200/200 requests without an OOM or request
+failure.
+
+| Implementation | Run 1 (req/s) | Run 2 (req/s) | Mean (req/s) |
+| --- | ---: | ---: | ---: |
+| `OmniKVTransferManager` | 4.4103 | 4.4497 | 4.4300 |
+| Native `MooncakeConnector` | 4.4457 | 4.4699 | 4.4578 |
+
+Native throughput is **0.63% higher** than the legacy mean. The absolute
+performance difference is below the 1% acceptance threshold.
+
+For accuracy, a non-zero transferred prefix and the resulting attention output
+were compared bit-for-bit with the legacy manager path:
+
+```text
+prefix_bits_equal    True
+attention_bits_equal True
+max_abs_diff         0.0
+```
+
+The one-AR/one-DiT TP=2 debug topology also completed 20/20 sequential
+requests. A separate real-TCP transport check transferred both CFG rows from a
+TP=1 producer to a strict-Ulysses SP=2 consumer with exact values on both
+target ranks.
+
+## Scope and limitations
+
+- The fully validated end-to-end topology is 2x AR TP=2 plus 2x DiT TP=2.
+- Heterogeneous TP=1 to SP=2 reshaping is transport-tested, but arbitrary mixed
+  TP/CFG/SP end-to-end combinations have not been exhaustively validated.
+- CFG-parallel end-to-end deployment is not part of this acceptance run.
+- Native AR-to-DiT transfer currently requires `async_chunk: false`.
+- TCP is covered by the acceptance run. The supplementary single-node RDMA
+  measurements use a separate short workload; they do not establish RDMA
+  support for multi-node deployment, which remains out of scope.
+- The native path currently validates the two-stage `AR(0) -> DiT(1)` graph. Other stage graphs are rejected during startup until
+  their producer/consumer edge mapping is implemented.
+- Sleep mode is disabled for the native paged path because registered KV pages
+  must remain mapped for the connector lifetime.
+
+- Ascend NPU adaptation and its performance validation are maintained in a
+  separate follow-up PR; the GPU acceptance results above do not validate NPU.
+- The homogeneous piecewise CUDA batching optimization is retained. A native
+  paged-attention CUDA graph capture/replay regression covers two CFG rows,
+  strided QKV, mixed causal/image spans, changed inputs, and output stability.
+  It must be run on CUDA hardware; CPU unit tests do not establish graph or
+  end-to-end performance acceptance.

tests/config/test_omni_config.py

@@ -62,6 +62,24 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 _DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy"
 
 
+@pytest.mark.parametrize("async_chunk", [False, True])
+def test_native_kv_transfer_requires_completed_ar_stage(async_chunk):
+    from types import SimpleNamespace
+
+    pipeline = SimpleNamespace(stages=(), model_type="test")
+    deploy = DeployConfig(
+        async_chunk=async_chunk,
+        stages=[
+            StageDeployConfig(stage_id=0, engine_extras={"kv_transfer_config": {"kv_connector": "MooncakeConnector"}})
+        ],
+    )
+    if async_chunk:
+        with pytest.raises(ValueError, match="requires async_chunk=False"):
+            omni_config_module._validate_async_chunk_support(pipeline, deploy)
+    else:
+        omni_config_module._validate_async_chunk_support(pipeline, deploy)
+
+
 @pytest.fixture(autouse=True)
 def _stable_test_platform(monkeypatch):
     from vllm_omni import platforms
@@ -921,6 +939,7 @@ def test_sub_config_fields_match_structured_scopes():
     assert {f.name for f in fields(OmniStageConnectorConfig)} == {
         "async_chunk",
         "omni_kv_config",
+        "kv_transfer_config",
         "stage_connector",
         "output_connectors",
         "input_connectors",

tests/core/sched/test_omni_ar_scheduler_free_request_cleanup.py

@@ -7,6 +7,7 @@ window fills with stale entries after K completions.
 
 from __future__ import annotations
 
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -15,6 +16,7 @@ import pytest
 # Request / RequestStatus are bound in this module. Ruff isort would reorder them.
 # isort: off
 import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
+from vllm.v1.request import RequestStatus
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 
 # isort: on
@@ -72,3 +74,85 @@ def test_free_request_is_safe_without_a_chunk_transfer_adapter():
     sched._free_request(request)  # must not raise
 
     sched._free_input_coordinator_request.assert_called_once_with("req-1")
+
+
+def test_native_mooncake_finish_uses_confirmed_token_boundary() -> None:
+    sched = _make_scheduler(chunk_transfer_adapter=None)
+    sched.vllm_config = SimpleNamespace(kv_transfer_config=SimpleNamespace(kv_connector="MooncakeConnector"))
+    observed = {}
+
+    def connector_finished(request):
+        observed["status"] = request.status
+        observed["num_computed_tokens"] = request.num_computed_tokens
+        return True, None
+
+    sched._connector_finished = connector_finished
+    request = _make_finished_request("req-native")
+    request.client_index = 0
+    request.status = RequestStatus.FINISHED_STOPPED
+    request.kv_transfer_params = {
+        "do_remote_decode": True,
+        "transfer_id": "xfer-req-native",
+    }
+    request.num_computed_tokens = 13
+    request.num_output_placeholders = 1
+    request.num_prompt_tokens = 10
+
+    kv_params, _ = sched._free_request(request)
+
+    assert observed == {
+        "status": RequestStatus.FINISHED_LENGTH_CAPPED,
+        "num_computed_tokens": 12,
+    }
+    assert request.status == RequestStatus.FINISHED_STOPPED
+    assert request.num_computed_tokens == 13
+    assert kv_params == {
+        "transfer_id": "xfer-req-native",
+        "num_transfer_tokens": 12,
+    }
+
+
+def test_native_mooncake_recycles_ar_pages_in_physical_order() -> None:
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched.vllm_config = SimpleNamespace(kv_transfer_config=SimpleNamespace(kv_connector="MooncakeConnector"))
+    blocks = [SimpleNamespace(block_id=7), SimpleNamespace(block_id=2), SimpleNamespace(block_id=4)]
+    sched.kv_cache_manager = SimpleNamespace(
+        enable_caching=False,
+        pop_blocks_for_free=MagicMock(return_value=blocks),
+        block_pool=SimpleNamespace(free_blocks=MagicMock()),
+    )
+    sched.defer_block_free = False
+    sched.processed_step_seq = 3
+    request = SimpleNamespace(last_sched_seq=3)
+
+    sched._free_request_blocks(request)
+
+    sched.kv_cache_manager.block_pool.free_blocks.assert_called_once()
+    freed = sched.kv_cache_manager.block_pool.free_blocks.call_args.args[0]
+    assert [block.block_id for block in freed] == [2, 4, 7]
+
+
+def test_native_mooncake_deferred_ar_pages_drain_in_physical_order() -> None:
+    from collections import deque
+
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched.vllm_config = SimpleNamespace(kv_transfer_config=SimpleNamespace(kv_connector="MooncakeConnector"))
+    blocks = [SimpleNamespace(block_id=7), SimpleNamespace(block_id=2), SimpleNamespace(block_id=4)]
+    sched.kv_cache_manager = SimpleNamespace(
+        enable_caching=False,
+        pop_blocks_for_free=MagicMock(return_value=blocks),
+        block_pool=SimpleNamespace(free_blocks=MagicMock()),
+    )
+    sched.defer_block_free = True
+    sched.processed_step_seq = 2
+    sched.sched_step_seq = 3
+    sched.deferred_frees = deque()
+    request = SimpleNamespace(last_sched_seq=3)
+
+    sched._free_request_blocks(request)
+    assert [block.block_id for block in sched.deferred_frees[0][1]] == [7, 4, 2]
+
+    sched.processed_step_seq = 3
+    sched._drain_deferred_frees()
+    freed = list(sched.kv_cache_manager.block_pool.free_blocks.call_args.args[0])
+    assert [block.block_id for block in freed] == [2, 4, 7]

tests/diffusion/diffusion_kv/test_block_tables.py

@@ -325,7 +325,7 @@ def _runner(
         model_config=SimpleNamespace(max_model_len=max_len),
         scheduler_config=SimpleNamespace(
             max_num_batched_tokens=max_num_batched_tokens,
-            max_num_seqs=max_num_seqs,
+            max_num_seqs=max_num_seqs * max_rows_per_request,
         ),
         parallel_config=SimpleNamespace(
             decode_context_parallel_size=cp_size,
@@ -403,7 +403,7 @@ def test_initialize_failure_does_not_publish_partial_cache_state(
     successful_init_cache = model_runner_backend_module.init_kv_cache
     vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(max_model_len=16),
-        scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=1),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=2),
         parallel_config=SimpleNamespace(
             decode_context_parallel_size=1,
             cp_kv_cache_interleave_size=1,
@@ -445,7 +445,7 @@ def test_initialize_failure_does_not_publish_partial_cache_state(
     assert backend.block_tables is None
     assert backend.paged_attention_adapter is None
     assert backend._kv_cache_layer_adapters["layer-0"].kv_cache == "placeholder-0"
-    assert backend.vllm_config.scheduler_config.max_num_seqs == 1
+    assert backend.vllm_config.scheduler_config.max_num_seqs == 2
 
     monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", successful_init_cache)
     backend.initialize_kv_cache(config)

tests/diffusion/diffusion_kv/test_initialization.py

@@ -58,6 +58,7 @@ def _od_config(**overrides):
         gpu_memory_utilization=0.9,
         kv_cache_memory_bytes=None,
         max_num_seqs=1,
+        diffusion_kv_max_rows_per_request=1,
         max_num_batched_tokens=64,
         num_gpus=1,
     )

tests/diffusion/diffusion_kv/test_kv_connector.py

@@ -1,16 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+from contextlib import nullcontext
+from types import SimpleNamespace
 from unittest import mock
 
 import pytest
 import torch
 from vllm.config import KVTransferConfig
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.outputs import KVConnectorOutput
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.kv_connector import (
+    KVTransferRegistrationError,
+    commit_kv_load,
     parse_kv_transfer_config,
     shutdown_kv_connector,
+    wait_for_kv_load,
 )
 from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler
 from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
@@ -138,3 +145,205 @@ def test_shutdown_kv_connector_is_idempotent() -> None:
 
     assert ensure_shutdown.call_count == 2
     scheduler_connector.shutdown.assert_called_once()
+
+
+def test_commit_aligns_cfg_transport_pages_without_extending_computed_prefix() -> None:
+    requests = tuple(
+        SimpleNamespace(
+            request_id=f"req/{sequence_id}",
+            kv_transfer_params={"num_transfer_tokens": 9},
+            num_tokens=16,
+            num_computed_tokens=0,
+        )
+        for sequence_id in range(2)
+    )
+    manager = SimpleNamespace(
+        get_blocks=lambda _request_id: KVCacheBlocks(([0, 1, 2, 3],)),
+        kv_cache_config=SimpleNamespace(kv_cache_groups=[SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=4))]),
+    )
+    connector = mock.Mock()
+
+    expected = commit_kv_load(connector, manager, requests, [7, 5])
+
+    assert expected == {"req/0", "req/1"}
+    calls = connector.update_state_after_alloc.call_args_list
+    assert [len(call.args[1].blocks[0]) for call in calls] == [3, 3]
+    assert [call.args[2] for call in calls] == [7, 5]
+    assert [request.num_computed_tokens for request in requests] == [7, 5]
+
+
+def test_wait_without_pending_load_uses_vllm_post_forward_once() -> None:
+    connector = mock.Mock()
+    output = mock.Mock(invalid_block_ids=set())
+    active_connector = mock.Mock(kv_connector=connector)
+    active_connector.post_forward.return_value = output
+    scheduler_output = SimpleNamespace(
+        kv_transfer_request_ids=set(),
+        finished_req_ids={"finished"},
+        kv_finished_request_ids={"finished/diffusion-kv/0", "finished/diffusion-kv/1"},
+    )
+
+    result = wait_for_kv_load(active_connector, scheduler_output, timeout=1.0)
+
+    assert result is output
+    active_connector.pre_forward.assert_called_once_with(scheduler_output)
+    active_connector.post_forward.assert_called_once_with({"finished/diffusion-kv/0", "finished/diffusion-kv/1"})
+    connector.get_finished.assert_not_called()
+
+
+def test_timeout_returns_partial_completion_without_releasing_pages():
+    active = mock.Mock()
+    active.kv_connector.get_finished.return_value = (set(), {"cfg0"})
+    active.kv_connector.get_block_ids_with_load_errors.return_value = set()
+    active.post_forward.return_value = KVConnectorOutput()
+    scheduled = SimpleNamespace(kv_transfer_request_ids={"cfg0", "cfg1"}, kv_finished_request_ids=set())
+    output = wait_for_kv_load(active, scheduled, timeout=0)
+    assert output.finished_recving == {"cfg0"}
+    active.kv_connector.request_finished.assert_not_called()
+    active.post_forward.assert_called_once_with(set())
+
+
+@pytest.fixture
+def cfg_registration():
+    requests = tuple(
+        SimpleNamespace(
+            request_id=f"cfg{i}",
+            status=None,
+            num_tokens=8,
+            kv_transfer_params={"num_transfer_tokens": 4, "transfer_id": "ticket", "do_remote_prefill": True},
+        )
+        for i in range(2)
+    )
+    manager = SimpleNamespace(
+        get_blocks=lambda _: KVCacheBlocks(([0, 1],)),
+        kv_cache_config=SimpleNamespace(kv_cache_groups=[SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=4))]),
+    )
+    return requests, manager
+
+
+def test_cfg_validation_is_atomic_before_connector_registration(cfg_registration):
+    requests, manager = cfg_registration
+    requests[1].kv_transfer_params["num_transfer_tokens"] = 9
+    connector = mock.Mock()
+    with pytest.raises(KVTransferRegistrationError, match="boundary"):
+        commit_kv_load(connector, manager, requests, [4, 4])
+    connector.update_state_after_alloc.assert_not_called()
+
+
+@pytest.mark.parametrize("fail_init", [False, True])
+def test_sp_connector_init_snapshots_topology_and_restores_model_tp(monkeypatch, fail_init):
+    import vllm.distributed.parallel_state as parallel_state
+
+    import vllm_omni.diffusion.diffusion_kv.kv_connector as module
+    import vllm_omni.diffusion.distributed.parallel_state as omni_parallel
+
+    tp = SimpleNamespace(world_size=1, rank_in_group=0)
+    sp = SimpleNamespace(world_size=2, rank_in_group=1)
+    monkeypatch.setattr(parallel_state, "_TP", tp)
+    monkeypatch.setattr(omni_parallel, "get_sp_group", lambda: sp)
+    snapshot: dict[str, int] = {}
+
+    def initialize(*args):
+        snapshot.update(
+            rank=parallel_state.get_tensor_model_parallel_rank(),
+            size=parallel_state.get_tensor_model_parallel_world_size(),
+        )
+        if fail_init:
+            raise RuntimeError("initialization failed")
+
+    monkeypatch.setattr(module, "ensure_kv_transfer_initialized", initialize)
+    config = SimpleNamespace(
+        kv_transfer_config=SimpleNamespace(kv_connector="MooncakeConnector", engine_id="dit"),
+        parallel_config=SimpleNamespace(prefill_context_parallel_size=2),
+    )
+    with pytest.raises(RuntimeError, match="initialization failed") if fail_init else nullcontext():
+        module.init_worker_kv_connector(config, object())
+    assert parallel_state._TP is tp
+    assert snapshot == {"rank": 1, "size": 2}
+
+
+def test_partial_mooncake_registration_is_replaced_by_empty_notifications(cfg_registration):
+    from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
+        MooncakeConnector,
+        MooncakeConnectorScheduler,
+    )
+
+    scheduler = object.__new__(MooncakeConnectorScheduler)
+    scheduler.is_kv_producer = False
+    scheduler.is_kv_consumer = True
+    scheduler._reqs_need_recv = {}
+    connector = object.__new__(MooncakeConnector)
+    connector.connector_scheduler = scheduler
+    requests, manager = cfg_registration
+    for request in requests:
+        request.kv_transfer_params.update(remote_engine_id="ar", remote_bootstrap_addr="http://localhost:8998")
+
+    def register(request, blocks, count):
+        scheduler._reqs_need_recv[request.request_id] = (request, blocks)
+        request.kv_transfer_params["do_remote_prefill"] = False
+        if request.request_id == "cfg1":
+            raise RuntimeError("partial registration")
+
+    connector.update_state_after_alloc = register
+    with pytest.raises(KVTransferRegistrationError, match="partial registration"):
+        commit_kv_load(connector, manager, requests, [4, 4])
+    assert set(scheduler._reqs_need_recv) == {"cfg0", "cfg1"}
+    assert all(blocks == [] for _, blocks in scheduler._reqs_need_recv.values())
+    metadata = scheduler.build_connector_meta(None)
+    # Even partial registration failure sends every row in one per-rank batch,
+    # preserving the producer's complete fan-out count for the shared ticket.
+    assert set(metadata.reqs_to_recv["ar"]) == {"cfg0", "cfg1"}
+    assert all(
+        meta.transfer_id == "ticket" and meta.local_block_ids == [] for meta in metadata.reqs_to_recv["ar"].values()
+    )
+
+
+@pytest.mark.parametrize("timeout", [0, -1, float("inf"), "60"])
+def test_native_transfer_timeout_must_be_positive_finite(timeout):
+    with pytest.raises(ValueError, match="transfer_timeout"):
+        parse_kv_transfer_config(
+            {
+                **KV_TRANSFER_CONFIG,
+                "kv_connector_extra_config": {"transfer_timeout": timeout},
+            }
+        )
+
+
+def test_receive_completion_requires_all_ranks_across_polls():
+    from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+    from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
+
+    executor = SimpleNamespace(od_config=SimpleNamespace(num_gpus=2))
+    executor.collective_rpc = mock.Mock(
+        side_effect=[
+            [KVConnectorOutput(finished_recving={"cfg0", "cfg1"}), KVConnectorOutput(finished_recving={"cfg0"})],
+            [KVConnectorOutput(), KVConnectorOutput(finished_recving={"cfg1"})],
+        ]
+    )
+    scheduled = 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,
+        kv_connector_metadata=object(),
+        kv_transfer_request_ids={"cfg0", "cfg1"},
+    )
+    result = DiffusionExecutor.prepare_kv_for_forward(executor, scheduled)
+    assert result.finished_recving == {"cfg0"}
+    scheduled.kv_transfer_request_ids = set()
+    result = DiffusionExecutor.prepare_kv_for_forward(executor, scheduled)
+    assert result.finished_recving == {"cfg1"}
+    assert not executor._kv_receive_completed_ranks
+
+
+def test_native_transfer_rejects_sleep_without_changing_legacy_sleep():
+    with pytest.raises(ValueError, match="registered pages must remain mapped"):
+        OmniDiffusionConfig.from_kwargs(
+            diffusion_kv_mode="paged_scheduler",
+            diffusion_kv_max_rows_per_request=2,
+            kv_transfer_config=dict(KV_TRANSFER_CONFIG),
+            enable_sleep_mode=True,
+        )
+    assert OmniDiffusionConfig.from_kwargs(enable_sleep_mode=True).enable_sleep_mode

tests/diffusion/diffusion_kv/test_metadata.py

@@ -61,13 +61,14 @@ class _MetadataScheduler(BaseScheduler):
 
 
 def test_scheduler_copies_opaque_params_to_sequence_requests() -> None:
-    params = {"transfer_id": "xfer-req-0", "do_remote_prefill": True}
+    params = {"transfer_id": "xfer-req-0", "do_remote_prefill": True, "num_transfer_tokens": 3}
     sequence = DiffusionKVRequest(
         "req-0/diffusion-kv/0",
         sequence_id=0,
         prefix_len=4,
         target_len=2,
         seq_len=6,
+        prompt_token_ids=[10, 11, 12, 13],
     )
     request = OmniDiffusionRequest(
         prompt="prompt",
@@ -83,6 +84,9 @@ def test_scheduler_copies_opaque_params_to_sequence_requests() -> None:
 
     assert state.diffusion_kv_requests[0].kv_transfer_params == params
     assert state.diffusion_kv_requests[0].kv_transfer_params is not params
+    assert sequence.prompt_token_ids == [10, 11, 12]
+    assert sequence.num_prompt_tokens == 3
+    assert sequence.seq_len == 6
 
 
 def test_diffusion_kv_metadata_uses_native_cache_group_block_ids() -> None:

tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise upstream Mooncake completion accounting without device/network I/O."""
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+
+import msgspec
+import pytest
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import mooncake_connector as mc
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.outputs import KVConnectorOutput
+
+from vllm_omni.diffusion.diffusion_kv.kv_connector import install_mooncake_cfg_fanout
+from vllm_omni.worker.mixins import OmniWorkerMixin
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture
+def producer(monkeypatch):
+    worker = object.__new__(mc.MooncakeConnectorWorker)
+    worker.shutdown = Mock()  # No hardware constructor, sockets or listener thread.
+    worker.is_kv_consumer = False
+    worker.tp_rank, worker.tp_size = 0, 1
+    worker.transfer_topo = SimpleNamespace(handshake_target_ranks=lambda size: list(range(size)))
+    worker.reqs_need_send, worker.finished_sending_reqs = {}, set()
+    for name in (
+        "kv_caches_base_addr",
+        "block_len_per_layer",
+        "kv_block_len_per_layer",
+        "registered_layer_names",
+        "registered_layer_indices",
+        "registered_group_indices",
+    ):
+        setattr(worker, name, [])
+    worker._get_transfer_regions = Mock(return_value=[])
+    worker._producer_cache_is_replicated = Mock(return_value=False)
+    worker._encoder = msgspec.msgpack.Encoder()
+    worker._sender_executor = None
+    worker._send_blocks = Mock(return_value=0)
+    worker.xfer_stats = Mock()
+    worker._build_transfer_params = AsyncMock(return_value=([1], [2], [16], [], None))
+    monkeypatch.setattr(mc, "_align_transfer_regions", lambda *_: ([], [], None))
+    monkeypatch.setattr(mc, "_validate_asymmetric_region_lengths", lambda **_: None)
+    monkeypatch.setattr(mc.envs, "VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", 0.05)
+    connector = object.__new__(mc.MooncakeConnector)
+    connector.connector_worker = worker
+    return connector
+
+
+async def register(worker, ticket):
+    worker.sender_loop = asyncio.get_running_loop()
+    for blocks in ([], [[10]]):
+        meta = mc.MooncakeConnectorMetadata()
+        meta.add_new_req(ticket, blocks, {"transfer_id": ticket}, load_remote_cache=False)
+        await worker.record_send_reqs(meta)
+    return worker.reqs_need_send[ticket]
+
+
+def packet(rank, ranks, tickets, empty=False):
+    return mc.MooncakeXferMetadata(
+        remote_hostname="localhost",
+        remote_port=1234 + rank,
+        remote_tp_size=ranks,
+        remote_tp_rank=rank,
+        req_blocks={
+            f"{ticket}/{row}": (ticket, [] if empty else [[100 + rank * 10 + row]])
+            for ticket, rows in tickets.items()
+            for row in range(rows)
+        },
+        kv_caches_base_addr=[],
+        block_lens=[],
+        kv_block_lens=[],
+    )
+
+
+async def send(worker, rank, ranks, tickets, empty=False):
+    sock = SimpleNamespace(send_multipart=AsyncMock())
+    await worker.send_kv_to_decode(str(rank).encode(), sock, packet(rank, ranks, tickets, empty))
+    return [
+        msgspec.msgpack.decode(c.args[0][1], type=mc.MooncakeXferResponse) for c in sock.send_multipart.call_args_list
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("rows", [1, 2])
+@pytest.mark.parametrize("ranks", [2, 4])
+@pytest.mark.parametrize("in_flight", [False, True])
+async def test_source_pages_wait_for_every_rank_and_row(producer, rows, ranks, in_flight):
+    worker = producer.connector_worker
+    entered, release = asyncio.Event(), asyncio.Event()
+    original = worker._build_transfer_params
+
+    async def plan(ready, meta, *regions):
+        if in_flight and meta.remote_tp_rank == ranks - 1:
+            entered.set()
+            await release.wait()
+        return await original(ready, meta, *regions)
+
+    worker._build_transfer_params = plan
+    install_mooncake_cfg_fanout(producer)
+    adapted = worker._build_transfer_params
+    install_mooncake_cfg_fanout(producer)
+    assert worker._build_transfer_params is adapted  # Do not multiply twice.
+    ticket = await register(worker, "request")
+    pending = asyncio.create_task(send(worker, ranks - 1, ranks, {"request": rows})) if in_flight else None
+    try:
+        if pending is not None:
+            await asyncio.wait_for(entered.wait(), 2)
+        for rank in range(ranks - 1):
+            await send(worker, rank, ranks, {"request": rows})
+            assert ticket.need_send == ranks * rows
+            assert ticket.sent == (rank + 1) * rows
+            assert worker.reqs_need_send["request"] is ticket
+            finished = await worker.fetch_finished_sending_reqs()
+            scheduler = SimpleNamespace(connector=None, requests={"request": object()}, _free_blocks=Mock())
+            Scheduler._update_from_kv_xfer_finished(scheduler, KVConnectorOutput(finished_sending=finished))
+            scheduler._free_blocks.assert_not_called()
+        release.set()
+        responses = (
+            await asyncio.wait_for(pending, 2) if pending else await send(worker, ranks - 1, ranks, {"request": rows})
+        )
+        assert sum(len(response.ok_reqs or []) for response in responses) == rows
+        assert all(not response.err_reqs for response in responses)
+        assert "request" not in worker.reqs_need_send
+        finished = await worker.fetch_finished_sending_reqs()
+        assert finished == {"request"}
+        Scheduler._update_from_kv_xfer_finished(scheduler, KVConnectorOutput(finished_sending=finished))
+        scheduler._free_blocks.assert_called_once()
+        assert not await worker.fetch_finished_sending_reqs()
+    finally:
+        release.set()
+        if pending is not None:
+            await asyncio.gather(pending, return_exceptions=True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("empty", [False, True])
+@pytest.mark.parametrize("ranks", [1, 2])
+async def test_mixed_tickets_and_empty_block_notifications(producer, empty, ranks):
+    worker = producer.connector_worker
+    if empty:
+        worker._build_transfer_params.return_value = ([], [], [], [], None)
+    install_mooncake_cfg_fanout(producer)
+    single, cfg = await register(worker, "single"), await register(worker, "cfg")
+    await asyncio.gather(*(send(worker, rank, ranks, {"single": 1, "cfg": 2}, empty) for rank in range(ranks)))
+    assert (single.need_send, single.sent) == (ranks, ranks)
+    assert (cfg.need_send, cfg.sent) == (2 * ranks, 2 * ranks)
+    assert await worker.fetch_finished_sending_reqs() == {"single", "cfg"}
+    assert not worker.reqs_need_send
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failure", ["row_count", "transport"])
+async def test_failed_rank_does_not_report_source_complete(producer, failure):
+    worker = producer.connector_worker
+    install_mooncake_cfg_fanout(producer)
+    ticket = await register(worker, "request")
+    await send(worker, 0, 2, {"request": 2})
+    if failure == "transport":
+        worker._send_blocks.return_value = -1
+    responses = await send(worker, 1, 2, {"request": 1 if failure == "row_count" else 2})
+    assert any(response.err_reqs for response in responses)
+    assert ticket.sent == 2 and ticket.need_send == 4 and ticket.sending == 0
+    assert worker.reqs_need_send["request"] is ticket
+    assert not await worker.fetch_finished_sending_reqs()
+    # Preserve upstream timeout cleanup after all in-flight writes have stopped.
+    ticket.expire_time = 0
+    assert await worker.fetch_finished_sending_reqs() == {"request"}
+    assert not worker.reqs_need_send
+
+
+def test_worker_installs_adapter_after_native_initialization(producer, monkeypatch):
+    import vllm.distributed.kv_transfer as kv_transfer
+
+    worker = producer.connector_worker
+    original = worker._build_transfer_params
+
+    class UpstreamWorker:
+        def initialize_from_config(self, config):
+            assert config == "config"
+            assert worker._build_transfer_params is original
+
+    class OmniWorker(OmniWorkerMixin, UpstreamWorker):
+        pass
+
+    monkeypatch.setattr(kv_transfer, "has_kv_transfer_group", lambda: True)
+    monkeypatch.setattr(kv_transfer, "get_kv_transfer_group", lambda: producer)
+    object.__new__(OmniWorker).initialize_from_config("config")
+    assert worker._build_transfer_params is not original
+
+
+@pytest.mark.parametrize("role", ["other_connector", "scheduler", "consumer"])
+def test_adapter_leaves_other_roles_untouched(producer, role):
+    worker = producer.connector_worker
+    original = worker._build_transfer_params
+    if role == "scheduler":
+        producer.connector_worker = None
+    elif role == "consumer":
+        worker.is_kv_consumer = True
+    install_mooncake_cfg_fanout(object() if role == "other_connector" else producer)
+    assert worker._build_transfer_params is original

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -94,7 +94,7 @@ class _FakeAttentionGroup:
         return self.builder
 
 
-@dataclass(frozen=True, slots=True)
+@dataclass(slots=True)
 class _FakeNativeMetadata:
     build_id: int
     causal: bool
@@ -102,6 +102,7 @@ class _FakeNativeMetadata:
     query_start_loc_cpu: torch.Tensor
     positions: torch.Tensor
     slot_mappings: torch.Tensor
+    max_num_splits: int = 0
 
 
 class _FakeLayer:
@@ -288,6 +289,8 @@ def test_omni_paged_backend_consumes_context_and_restores_diffusion_shape(
     assert output.shape == query.shape
     assert torch.equal(output, query)
     assert layer.calls[0][0].shape == (5, 2, 4)
+    assert layer.calls[0][1] is None
+    assert layer.calls[0][2] is None
     assert layer.native_events == ["update", "forward"]
     assert layer.calls[0][3] is events[0][2]
 
@@ -440,9 +443,7 @@ def test_platform_prewrite_updates_once_then_piecewise_reads_cache(monkeypatch:
 
 def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, layer, events = _make_adapter(monkeypatch)
-    # Exercise the Ascend-only homogeneous batch path. CUDA keeps its native
-    # output-buffer contract and is covered by the GPU adapter tests.
-    monkeypatch.setattr(adapter_module.current_omni_platform, "is_npu", lambda: True)
+    layer.impl.vllm_flash_attn_version = 2
     batch = adapter.prepare_batch(
         [
             DiffusionPagedAttentionRow(
@@ -499,10 +500,10 @@ def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.
     assert torch.equal(layer.updates[0][0], key.reshape(12, 2, 4))
     assert torch.equal(layer.updates[0][1], value.reshape(12, 2, 4))
     assert [call[0].shape[0] for call in layer.calls] == [4, 6, 2]
-    assert [call[1].shape[0] for call in layer.calls] == [4, 6, 2]
-    assert [call[2].shape[0] for call in layer.calls] == [4, 6, 2]
+    assert all(call[1] is None and call[2] is None for call in layer.calls)
     assert all(call[3] is event[2] for call, event in zip(layer.calls, events[1:], strict=True))
     assert [metadata.build_id for metadata in context.piecewise_native_metadata] == [1, 2, 3]
+    assert [metadata.max_num_splits for metadata in context.piecewise_native_metadata] == [1, 1, 1]
 
     # The first metadata build is the normal whole-query path. Piecewise calls
     # then use causal [3, 5), full [5, 8), and causal [8, 9) segments.
@@ -539,6 +540,51 @@ def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.
     ]
 
 
+def test_homogeneous_piecewise_packs_strided_query_only_per_segment(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    batch = adapter.prepare_batch(
+        [
+            DiffusionPagedAttentionRow(
+                request_id="req-0",
+                sequence_id=0,
+                query_len=6,
+                seq_len=9,
+                kv_start_pos=3,
+            ),
+            DiffusionPagedAttentionRow(
+                request_id="req-1",
+                sequence_id=1,
+                query_len=6,
+                seq_len=9,
+                kv_start_pos=3,
+            ),
+        ]
+    )
+    packed_qkv = torch.randn(2, 6, 3, 2, 4)
+    query, key, value = packed_qkv.unbind(dim=2)
+    assert not query.is_contiguous()
+    metadata = SimpleNamespace(
+        attn_mask=None,
+        full_attn_spans=[[(5, 8)], [(5, 8)]],
+        query_ranges=None,
+        extra={},
+    )
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context(
+            "layer-0",
+            query,
+            key,
+            value,
+            omni_attn_metadata=metadata,
+        )
+
+    assert not context.query.is_contiguous()
+    assert context.query.data_ptr() == query.data_ptr()
+
+
 def test_piecewise_metadata_snapshots_reused_native_scheduler_buffer(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, _, _ = _make_adapter(monkeypatch)
     shared_scheduler_metadata = torch.zeros(1, dtype=torch.int32)
@@ -1232,7 +1278,7 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         ulysses_degree=2,
     )
 
-    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 2)]
+    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 1)]
     assert specialized_backends == [(native_backend, 2)]
     assert config.attention_config.backend is None
     assert config.attention_config.backend_per_kind is original_backend_per_kind

tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py

@@ -43,8 +43,7 @@ class _SmokeDiffusionAttention(nn.Module):
         self.softmax_scale = _HEAD_SIZE**-0.5
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
-def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> None:
+def _native_adapter(num_rows: int = 1):
     device = torch.device("cuda", torch.accelerator.current_device_index())
     vllm_config = VllmConfig(
         cache_config=CacheConfig(
@@ -89,8 +88,8 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
 
     canonical_spec = native_layer.spec
     kv_cache_config = KVCacheConfig(
-        num_blocks=4,
-        kv_cache_tensors=[build_kv_cache_tensor(canonical_spec, 4, [_LAYER_NAME])],
+        num_blocks=4 * num_rows,
+        kv_cache_tensors=[build_kv_cache_tensor(canonical_spec, 4 * num_rows, [_LAYER_NAME])],
         kv_cache_groups=[
             KVCacheGroupSpec(
                 layer_names=[_LAYER_NAME],
@@ -120,14 +119,15 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
         )
         block_tables = BlockTables(
             block_sizes=[_BLOCK_SIZE],
-            max_num_reqs=1,
-            max_num_batched_tokens=32,
+            max_num_reqs=num_rows,
+            max_num_batched_tokens=32 * num_rows,
             max_num_blocks_per_group=[2],
             device=device,
             kernel_block_sizes=kernel_block_sizes,
         )
 
-    block_tables.append_block_ids(0, ([3, 1],), overwrite=True)
+    for row in range(num_rows):
+        block_tables.append_block_ids(row, ([4 * row + 3, 4 * row + 1],), overwrite=True)
     block_tables.apply_staged_writes()
     adapter = DiffusionPagedAttentionAdapter(
         vllm_config=vllm_config,
@@ -136,11 +136,17 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
         block_tables=block_tables,
         attn_groups=attn_groups,
         layers={_LAYER_NAME: native_layer},
-        resolve_row=lambda _request_id, _sequence_id, _context_id: DiffusionPagedAttentionRowBinding(
-            row_index=0,
+        resolve_row=lambda _request_id, sequence_id, _context_id: DiffusionPagedAttentionRowBinding(
+            row_index=sequence_id,
             max_seq_len=19,
         ),
     )
+    return adapter, diffusion_layer, device
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
+def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> None:
+    adapter, diffusion_layer, device = _native_adapter()
     prefix_batch = adapter.prepare_batch(
         [
             DiffusionPagedAttentionRow(
@@ -258,3 +264,48 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
     torch.testing.assert_close(prefix_output, prefix_reference, rtol=2e-2, atol=2e-2)
     torch.testing.assert_close(suffix_output, suffix_reference, rtol=2e-2, atol=2e-2)
     torch.testing.assert_close(piecewise_output, piecewise_reference, rtol=2e-2, atol=2e-2)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
+def test_hunyuan_homogeneous_cfg_paged_attention_graph_replay() -> None:
+    """Real paged kernels, strided QKV, mixed image spans and changed replay inputs."""
+    adapter, layer, device = _native_adapter(num_rows=2)
+    packed = torch.randn(2, 19, 3, _NUM_HEADS, _HEAD_SIZE, device=device, dtype=torch.float16)
+    query, key, value = packed.unbind(dim=2)
+    backend = FlashAttentionImpl(
+        num_heads=_NUM_HEADS, head_size=_HEAD_SIZE, softmax_scale=layer.softmax_scale, num_kv_heads=_NUM_HEADS
+    )
+    batch = adapter.prepare_batch(
+        [DiffusionPagedAttentionRow(request_id="cfg", sequence_id=i, query_len=19, seq_len=19) for i in range(2)]
+    )
+    metadata = AttentionMetadata(full_attn_spans=[[(5, 10)], [(5, 10)]])
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context(_LAYER_NAME, query, key, value, omni_attn_metadata=metadata)
+    assert context.piecewise_plan.homogeneous_batch_shape == (2, 19)
+    warmup = torch.cuda.Stream()
+    warmup.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(warmup):
+        for _ in range(3):
+            backend.forward_paged(context)
+    torch.cuda.current_stream().wait_stream(warmup)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = backend.forward_paged(context)
+    output_ptr = captured.data_ptr()
+    mask = torch.ones(19, 19, device=device, dtype=torch.bool).tril()
+    mask[5:10, 5:10] = True
+    for _ in range(2):
+        packed.copy_(torch.randn_like(packed))
+        graph.replay()
+        actual = captured.clone()
+        eager = backend.forward_paged(context)
+        reference = F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            attn_mask=mask,
+            scale=layer.softmax_scale,
+        ).transpose(1, 2)
+        torch.testing.assert_close(actual.reshape_as(reference), reference, rtol=2e-2, atol=2e-2)
+        torch.testing.assert_close(actual, eager, rtol=2e-2, atol=2e-2)
+        assert captured.data_ptr() == output_ptr

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -176,10 +176,11 @@ def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None:
     profile_request = object()
 
     assert worker.determine_available_kv_memory(profile_request) == [4096]
-    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+    worker.model_runner.profile_run.assert_not_called()
 
 
 def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch) -> None:
+    monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", SimpleNamespace(empty_cache=Mock()))
     worker = object.__new__(DiffusionWorker)
     worker.rank = 0
     worker.init_snapshot = SimpleNamespace(free_memory=900)
@@ -209,6 +210,7 @@ def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch
 
 @pytest.mark.parametrize("non_kv_cache_memory", [750, 800])
 def test_worker_rejects_non_positive_profiled_kv_memory(monkeypatch, non_kv_cache_memory: int) -> None:
+    monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", SimpleNamespace(empty_cache=Mock()))
     worker = object.__new__(DiffusionWorker)
     worker.rank = 0
     worker.init_snapshot = SimpleNamespace(free_memory=900)
@@ -245,6 +247,7 @@ def test_worker_rejects_non_positive_profiled_kv_memory(monkeypatch, non_kv_cach
 
 
 def test_worker_profiles_activation_headroom_instead_of_current_residency(monkeypatch) -> None:
+    monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", SimpleNamespace(empty_cache=Mock()))
     worker = object.__new__(DiffusionWorker)
     worker.rank = 0
     worker.init_snapshot = SimpleNamespace(free_memory=1200)
@@ -310,3 +313,31 @@ def test_rank_probe_gathers_local_failure_before_raising(monkeypatch) -> None:
         )
 
     assert gathered == [(False, "ValueError: local failure")]
+
+
+def test_hunyuan_native_layer_identity_matches_ar_spec(monkeypatch):
+    import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer as hy3
+
+    monkeypatch.setattr(hy3, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(hy3, "get_sequence_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(hy3, "get_allgather_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(hy3, "get_sequence_parallel_rank", lambda: 0)
+    monkeypatch.setattr(hy3, "QKVParallelLinear", lambda **_: nn.Identity())
+    monkeypatch.setattr(hy3, "RowParallelLinear", lambda **_: nn.Identity())
+    monkeypatch.setattr(hy3, "get_rope", lambda **_: nn.Identity())
+
+    # Only backend construction is fake; execute both HY3 constructors and the
+    # real runner spec-discovery path, including the actual Attention prefix.
+    def attention(**kwargs):
+        return _attention(enabled=kwargs.get("paged_kv_cache_role") is not None, prefix=kwargs["prefix"])
+
+    monkeypatch.setattr(hy3, "Attention", attention)
+    model = hy3.HunYuanAttention(
+        config=SimpleNamespace(num_key_value_heads=2, attention_head_dim=8),
+        hidden_size=16,
+        num_heads=2,
+        num_kv_heads=2,
+        prefix="layers.3.self_attn",
+    )
+    runner = _runner(model.image_attn.attn)
+    assert set(runner.get_kv_cache_spec()) == {"model.layers.3.self_attn.attn"}

tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.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 types import SimpleNamespace
 
@@ -362,6 +362,20 @@ def test_prepared_model_inputs_match_local_tokenization(
     local_inputs = pipeline.prepare_model_inputs(**common_kwargs)
     prepared_inputs = pipeline.prepare_model_inputs(**common_kwargs, prepared_layout=prepared_layout)
 
+    if with_reference_image:
+        # Preparation runs without a forward context. Coverage must come from
+        # the scheduler metadata copied onto this request, including every CFG row.
+        reused = pipeline.prepare_model_inputs(
+            **common_kwargs, prepared_layout=prepared_layout, kv_computed_tokens=(5,) * len(prefix_lens)
+        )
+        assert reused["cond_vae_images"] is None and reused["cond_vit_images"] is None
+        partial = pipeline.prepare_model_inputs(
+            **common_kwargs,
+            prepared_layout=prepared_layout,
+            kv_computed_tokens=(5,) * (len(prefix_lens) - 1) + (4,),
+        )
+        assert partial["cond_vae_images"] is not None and partial["cond_vit_images"] is not None
+
     comparable_fields = (
         "input_ids",
         "position_ids",
@@ -474,3 +488,33 @@ def test_paged_preprocess_attaches_layout_and_scheduler_kv_requests(monkeypatch)
         (12, 17, 32)
     ]
     assert len(tokenizer.calls) == 1
+
+
+@pytest.mark.parametrize("transfer_tokens, expected", [(9, [7, 5]), (4, [4, 4])])
+def test_cfg_rows_preserve_reusable_token_ids_and_native_transfer_boundary(transfer_tokens, expected):
+    from vllm_omni.diffusion.diffusion_kv.kv_connector import prepare_kv_requests
+
+    tokenizer, image_processor = _components([12, 14])
+    request = _request(guidance_scale=5.0)
+    layout = _prepare(request, tokenizer, image_processor)
+    layout.tokenizer_output.think_recaption_end_pos = [[7], [5]]
+    layout.tokenizer_output.uncond_cfg_start_pos = [[None], [5]]
+    rows = build_hunyuan_diffusion_kv_requests(request, layout)
+    assert [row.prompt_token_ids for row in rows] == [
+        layout.tokenizer_output.tokens[0, :7].tolist(),
+        layout.tokenizer_output.tokens[1, :5].tolist(),
+    ]
+    prepare_kv_requests(rows, {"num_transfer_tokens": transfer_tokens})
+    assert [row.num_prompt_tokens for row in rows] == expected
+    assert all(row.kv_transfer_params["num_transfer_tokens"] == transfer_tokens for row in rows)
+    layout.tokenizer_output.uncond_cfg_start_pos = [[5]]
+    with pytest.raises(ValueError, match="one boundary per CFG row"):
+        build_hunyuan_diffusion_kv_requests(request, layout)
+
+
+@pytest.mark.parametrize("computed, expected", [((8, 8), True), ((8, 5), False), ((0, 0), False), ((), False)])
+def test_native_image_reuse_requires_coverage_in_every_cfg_row(computed, expected):
+    from vllm_omni.diffusion.models.hunyuan_image3.request_layout import native_kv_covers_cond_images
+
+    output = TokenizerEncodeOutput(joint_image_slices=[[slice(2, 8)], [slice(2, 6)]])
+    assert native_kv_covers_cond_images(output, computed) is expected

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 import math
 from contextlib import ExitStack, contextmanager
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
@@ -143,6 +144,104 @@ def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
     assert mgr.state_dict() == {}
 
 
+def test_cache_manager_initializes_with_quantization_skip_layers(monkeypatch):
+    import vllm_omni.diffusion.attention.layer as layer
+    from vllm_omni.diffusion.config import set_current_diffusion_config
+    from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import ImageKVCacheManager
+
+    config = SimpleNamespace(
+        diffusion_attention_config=None,
+        parallel_config=SimpleNamespace(ring_degree=1),
+        diffusion_kv_cache_dtype="auto",
+        diffusion_kv_cache_skip_layer_indices={3},
+    )
+    monkeypatch.setattr(layer, "get_attn_backend_for_role", lambda **_: (layer.SDPABackend, None))
+    monkeypatch.setattr(layer, "build_parallel_attention_strategy", lambda **_: object())
+    with patched_mgr_env(), set_current_diffusion_config(config):
+        # Keep the real Attention constructor and its skip-layer validation.
+        with patch(f"{_TRANSFORMER_MODULE}.Attention", layer.Attention):
+            mgr = ImageKVCacheManager(
+                num_heads=NUM_HEADS,
+                num_kv_heads=NUM_KV_HEADS,
+                head_dim=HEAD_DIM,
+                scaling=SCALING,
+                prefix="model.layers.3.self_attn",
+            )
+    assert mgr.attn.layer_idx == 3
+    assert not mgr.attn._should_apply_kv_cache_quant()
+
+
+@pytest.mark.parametrize("first_step", [True, False])
+def test_native_connector_keeps_imported_prefix_in_paged_attention(first_step) -> None:
+    from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
+
+    mgr = _make_cache_mgr()
+    mgr.attn.paged_kv_active = True
+    query_len, prefix_len = IMAGE_TOKEN_LEN, 19
+    key, value = _make_known_kv(2 * query_len)
+    spans = [[(prefix_len, prefix_len + query_len)]] * 2
+    with (
+        set_forward_context(),
+        override_paged_kv_adapter(object()),
+        patch.object(mgr, "_forward_dense_legacy", side_effect=AssertionError("dense KV path")),
+        patch(f"{_TRANSFORMER_MODULE}.repeat_kv", side_effect=AssertionError("KV head expansion")),
+        patch(f"{_TRANSFORMER_MODULE}.get_paged_kv_computed_tokens", return_value=(prefix_len, prefix_len)),
+    ):
+        _call_mgr(
+            mgr,
+            2,
+            query_len,
+            prefix_len + query_len,
+            key,
+            value,
+            first_step=first_step,
+            full_attn_spans=spans,
+        )
+    assert len(mgr.attn.paged_calls) == 1
+    assert mgr.attn.paged_calls[0][1].shape == (2, query_len, NUM_KV_HEADS, HEAD_DIM)
+    assert mgr.attn.paged_metadata[0].attn_mask is None
+    assert mgr.attn.paged_metadata[0].full_attn_spans == spans
+    assert mgr.image_kv_cache_map is None
+    assert not hasattr(mgr, "_materialize_paged_prefix")
+
+
+def test_legacy_tensor_prefix_still_uses_dense_cache_and_reuse() -> None:
+    mgr = _make_cache_mgr()
+    prefix_len, prompt_len = 3, 2
+    imported = _make_known_kv(prefix_len, base=10.0)
+    mgr._injected_ar_kv = [imported]
+    first_len = prompt_len + IMAGE_TOKEN_LEN
+    key, value = _make_known_kv(first_len, base=20.0)
+    with (
+        patch.object(mgr, "_forward_paged", side_effect=AssertionError("legacy entered paged path")),
+    ):
+        _call_mgr(
+            mgr,
+            1,
+            first_len,
+            prefix_len + first_len,
+            key,
+            value,
+            first_step=True,
+            gen_timestep_scatter_index=_gen_timestep_index(1, prompt_len),
+        )
+        assert mgr.image_kv_cache_map is not None
+        assert mgr.image_kv_cache_map[0].shape[1] == prefix_len + prompt_len
+        next_key, next_value = _make_known_kv(IMAGE_TOKEN_LEN, base=30.0)
+        _call_mgr(
+            mgr,
+            1,
+            IMAGE_TOKEN_LEN,
+            prefix_len + first_len,
+            next_key,
+            next_value,
+            position_ids=torch.arange(prefix_len + prompt_len, prefix_len + first_len).unsqueeze(0),
+        )
+    assert len(mgr.attn.calls) == 2
+    assert not mgr.attn.paged_calls
+    assert mgr.attn.calls[-1][1].shape == (1, prefix_len + first_len, NUM_HEADS, HEAD_DIM)
+
+
 def test_scheduler_paged_kv_runs_piecewise_for_first_and_later_steps() -> None:
     from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
 

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -474,12 +474,13 @@ async def test_cancelled_future_does_not_kill_busy_loop():
 
 
 @pytest.mark.asyncio
-async def test_native_kv_reservation_error_wakes_stream_without_killing_busy_loop():
+@pytest.mark.parametrize("error_type", [ValueError, RuntimeError])
+async def test_native_kv_reservation_error_wakes_stream_without_killing_busy_loop(error_type):
     loop = asyncio.get_running_loop()
     engine = _make_engine_with_loop(loop)
 
     def fail_reservation(*args, **kwargs):
-        raise RuntimeError("native allocation bug")
+        raise error_type("native allocation bug")
 
     engine.scheduler._diffusion_kv_manager = SimpleNamespace(
         has_request=lambda request_id: False,

tests/diffusion/test_diffusion_model_runner.py

@@ -431,7 +431,13 @@ def test_update_states_carries_prepared_layout() -> None:
     request.prepared_layout = prepared_layout
     scheduler_output = SimpleNamespace(
         finished_req_ids=set(),
-        scheduled_new_reqs=[SimpleNamespace(request_id=request.request_id, req=request)],
+        scheduled_new_reqs=[
+            SimpleNamespace(
+                request_id=request.request_id,
+                req=request,
+                diffusion_kv_metadata=None,
+            )
+        ],
         scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
     )
 

tests/diffusion/test_diffusion_scheduler.py

@@ -11,6 +11,7 @@ import torch
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
 from pytest_mock import MockerFixture
 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
+from vllm.v1.outputs import KVConnectorOutput
 
 from tests.helpers.kv_layout import build_kv_cache_tensor
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
@@ -136,6 +137,55 @@ def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) ->
     )
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("scheduler_cls", [RequestScheduler, StepScheduler])
+@pytest.mark.parametrize("failure", ["timeout", "registration"])
+async def test_single_native_kv_failure_reaches_output_stream(mocker, scheduler_cls, failure):
+    from vllm_omni.diffusion.diffusion_kv.kv_connector import KVTransferRegistrationError
+
+    scheduler = scheduler_cls()
+    _initialize_paged_scheduler(scheduler, max_num_seqs=1)
+    connector = mocker.Mock()
+    connector.get_num_new_matched_tokens.return_value = (4, True)
+    connector.request_finished.return_value = (False, None)
+    scheduler._kv_connector = connector
+    request = _make_request("failed")
+    _attach_diffusion_kv(request)
+    request.diffusion_kv_requests[0].prompt_token_ids = [1, 2, 3, 4]
+    request.kv_transfer_params = {"num_transfer_tokens": 4, "do_remote_prefill": True}
+    scheduler.add_request(request)
+    if failure == "registration":
+        mocker.patch(
+            "vllm_omni.diffusion.sched.base_scheduler.commit_kv_load",
+            side_effect=KVTransferRegistrationError("registration failed"),
+        )
+    engine = object.__new__(DiffusionEngine)
+    engine.scheduler = scheduler
+    engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
+    engine.execution_mode = (
+        DiffusionExecutionMode.STEP_BATCH if scheduler_cls is StepScheduler else DiffusionExecutionMode.REQUEST_BATCH
+    )
+    engine.abort_queue = queue.Queue()
+    engine._cv = threading.Condition()
+    engine.main_loop = asyncio.get_running_loop()
+    stream = asyncio.Queue()
+    engine._out_streams = {"failed": stream}
+    engine.executor = mocker.Mock()
+    engine.executor.prepare_kv_for_forward.return_value = KVConnectorOutput()
+    scheduled = scheduler.schedule()
+    engine._prepare_kv_for_forward(scheduled)
+    assert scheduled.scheduled_request_ids == []
+    output = BatchRunnerOutput.from_list([])
+    finished = scheduler.update_from_output(scheduled, output)
+    engine._emit_outputs(finished, scheduled.scheduled_request_ids, output)
+    terminal = await asyncio.wait_for(stream.get(), timeout=1)
+    assert terminal.finished
+    assert terminal.error == ("Timed out receiving diffusion KV" if failure == "timeout" else "registration failed")
+    assert scheduler.get_request_state("failed") is None
+    # Reporting failure must not recycle pages that the sender can still write.
+    assert scheduler._diffusion_kv_manager.has_request("failed") == (failure == "timeout")
+
+
 class _StubScheduler:
     def __init__(self, request: OmniDiffusionRequest, output) -> None:
         self._request = request
@@ -602,16 +652,22 @@ class TestRequestScheduler:
         assert sched_output.finished_req_ids == {"impossible"}
         assert _new_ids(sched_output) == ["schedulable"]
 
-    def test_diffusion_kv_internal_allocation_error_finishes_request(self, monkeypatch) -> None:
+    def test_diffusion_kv_internal_allocation_error_is_request_scoped(self, monkeypatch) -> None:
         _initialize_paged_scheduler(self.scheduler)
         request = _make_request("native-error")
         _attach_diffusion_kv(request)
         self.scheduler.add_request(request)
         manager = self.scheduler._diffusion_kv_manager
         assert manager is not None
+        reserve_request = manager.reserve_request
+        next_request = _make_request("after-error")
+        _attach_diffusion_kv(next_request)
+        self.scheduler.add_request(next_request)
 
         def raise_native_error(*args, **kwargs):
-            raise ValueError("native allocation bug")
+            if args[0] == "native-error":
+                raise ValueError("native allocation bug")
+            return reserve_request(*args, **kwargs)
 
         monkeypatch.setattr(
             manager,
@@ -619,14 +675,150 @@ class TestRequestScheduler:
             raise_native_error,
         )
 
-        sched_output = self.scheduler.schedule()
-
+        output = self.scheduler.schedule()
         state = self.scheduler.get_request_state("native-error")
-        assert state is not None
         assert state.status == DiffusionRequestStatus.FINISHED_ERROR
         assert state.error == "native allocation bug"
-        assert sched_output.finished_req_ids == {"native-error"}
-        assert sched_output.scheduled_request_ids == []
+        assert output.finished_req_ids == {"native-error"}
+        assert _new_ids(output) == ["after-error"]
+        assert not manager.has_request("native-error")
+
+    @pytest.mark.parametrize("action", ["abort", "timeout", "preempt"])
+    def test_diffusion_kv_loading_defers_free_and_blocks_preempt(self, mocker: MockerFixture, action: str) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=5)
+        connector = mocker.Mock()
+        connector.get_num_new_matched_tokens.return_value = (4, True)
+        connector.request_finished.return_value = (False, None)
+        self.scheduler._kv_connector = connector
+        request = _make_request("loading")
+        request.diffusion_kv_requests = tuple(
+            DiffusionKVRequest(
+                f"loading/diffusion-kv/{i}",
+                sequence_id=i,
+                prefix_len=4,
+                target_len=4,
+                seq_len=8,
+                prompt_token_ids=[1, 2, 3, 4],
+            )
+            for i in range(2)
+        )
+        request.kv_transfer_params = {"num_transfer_tokens": 4, "do_remote_prefill": True}
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+        pool = manager.native_manager.block_pool
+        initial_free_blocks = pool.get_num_free_blocks()
+        self.scheduler.add_request(request)
+        scheduled = self.scheduler.schedule()
+        internal_ids = {f"loading/diffusion-kv/{i}" for i in range(2)}
+        assert scheduled.kv_transfer_request_ids == internal_ids
+        assert pool.get_num_free_blocks() == initial_free_blocks - 4
+        metadata = manager.get_metadata("loading")
+        state = self.scheduler.get_request_state("loading")
+        free_request = mocker.spy(manager, "free_request")
+
+        if action == "abort":
+            self.scheduler.finish_requests("loading", DiffusionRequestStatus.FINISHED_ABORTED)
+        elif action == "timeout":
+            assert self.scheduler.fail_incomplete_kv_loads(internal_ids) == {"loading"}
+            assert state.error == "Timed out receiving diffusion KV"
+        if action != "preempt":
+            assert state.is_finished()
+            assert self.scheduler.pop_request_state("loading") is state
+            with pytest.raises(ValueError, match="already active"):
+                self.scheduler.add_request(request)
+            assert self.scheduler.get_diffusion_kv_cleanup_targets(["loading"]) == []
+
+        # Empty, unrelated, and partial completion retain both CFG rows even
+        # after the frontend has consumed and popped the terminal request.
+        for finished in (set(), {"other/diffusion-kv/0"}, {"loading/diffusion-kv/0"}):
+            self.scheduler.update_kv_connector_output(KVConnectorOutput(finished_recving=finished))
+            assert "loading" in self.scheduler._kv_loading_request_ids
+            assert self.scheduler.preempt_request("loading") is False
+            assert not self.scheduler.completed_kv_drains()
+            assert manager.has_request("loading")
+            assert manager.get_metadata("loading") == metadata
+            assert pool.get_num_free_blocks() == initial_free_blocks - 4
+            free_request.assert_not_called()
+            connector.request_finished.assert_not_called()
+
+        # CFG completion may arrive in different polls, after rank aggregation.
+        self.scheduler.update_kv_connector_output(KVConnectorOutput(finished_recving={"loading/diffusion-kv/1"}))
+        assert "loading" not in self.scheduler._kv_loading_request_ids
+        if action == "preempt":
+            assert self.scheduler.preempt_request("loading") is True
+            assert state.status == DiffusionRequestStatus.PREEMPTED
+            assert manager.get_metadata("loading") == metadata
+            self.scheduler.finish_requests("loading", DiffusionRequestStatus.FINISHED_ABORTED)
+        else:
+            assert self.scheduler.completed_kv_drains() == {"loading"}
+            assert self.scheduler.get_diffusion_kv_cleanup_targets(["loading"]) == [
+                ("loading", metadata.allocation_generation)
+            ]
+            self.scheduler.release_kv_drains({"loading"})
+        assert not self.scheduler._running and not self.scheduler._waiting
+        assert self.scheduler._kv_finished_request_ids == internal_ids
+        assert connector.request_finished.call_count == 2
+        free_request.assert_called_once_with("loading")
+        assert not manager.has_request("loading")
+        assert pool.get_num_free_blocks() == initial_free_blocks
+
+    @pytest.mark.parametrize("cancel", [False, True])
+    def test_engine_timeout_keeps_serving_and_reclaims_late_receive(self, mocker, cancel):
+        _initialize_paged_scheduler(self.scheduler, num_blocks=16, max_num_seqs=2)
+        connector = mocker.Mock()
+        connector.get_num_new_matched_tokens.return_value = (4, True)
+        connector.request_finished.return_value = (False, None)
+        self.scheduler._kv_connector = connector
+        for rid in ("slow", "ready"):
+            req = _make_request(rid)
+            _attach_diffusion_kv(req)
+            req.diffusion_kv_requests[0].prompt_token_ids = [1, 2, 3, 4]
+            req.kv_transfer_params = {"num_transfer_tokens": 4, "do_remote_prefill": True}
+            self.scheduler.add_request(req)
+        scheduled = self.scheduler.schedule()
+        engine = object.__new__(DiffusionEngine)
+        engine.scheduler = self.scheduler
+        engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
+        engine.abort_queue = queue.Queue()
+        engine._cv = threading.Condition()
+        if cancel:
+            engine.abort_queue.put("slow")
+        engine.executor = mocker.Mock()
+        engine.executor.prepare_kv_for_forward.side_effect = [
+            KVConnectorOutput(finished_recving={"ready/diffusion-kv/0"}),
+            KVConnectorOutput(finished_recving={"slow/diffusion-kv/0"}),
+        ]
+        fail_engine = mocker.patch.object(engine, "_fail_engine")
+        engine._prepare_kv_for_forward(scheduled)
+        assert scheduled.scheduled_request_ids == ["ready"]
+        assert scheduled.finished_req_ids == {"slow"}
+        expected_status = DiffusionRequestStatus.FINISHED_ABORTED if cancel else DiffusionRequestStatus.FINISHED_ERROR
+        assert self.scheduler.get_request_state("slow").status == expected_status
+        assert self.scheduler._diffusion_kv_manager.has_request("slow")
+        engine._remove_diffusion_kv_requests(["slow"])
+        engine.executor.remove_diffusion_kv_requests.assert_not_called()
+        self.scheduler.pop_request_state("slow")
+        generation = self.scheduler._kv_request_generations["slow"]
+        engine._prepare_kv_for_forward(self.scheduler.schedule())
+        engine.executor.remove_diffusion_kv_requests.assert_called_once_with([("slow", generation)])
+        assert not self.scheduler._diffusion_kv_manager.has_request("slow")
+        assert self.scheduler._diffusion_kv_manager.has_request("ready")
+        fail_engine.assert_not_called()
+
+    def test_slow_kv_match_does_not_block_later_request(self, mocker):
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2)
+        connector = mocker.Mock()
+        connector.get_num_new_matched_tokens.side_effect = lambda req, _: (
+            (None, False) if req.request_id.startswith("slow/") else (0, False)
+        )
+        self.scheduler._kv_connector = connector
+        for rid in ("slow", "ready"):
+            request = _make_request(rid)
+            _attach_diffusion_kv(request)
+            self.scheduler.add_request(request)
+        output = self.scheduler.schedule()
+        assert _new_ids(output) == ["ready"]
+        assert list(self.scheduler._waiting) == ["slow"]
 
     def test_diffusion_kv_preemption_retains_allocation(self) -> None:
         _initialize_paged_scheduler(self.scheduler, num_blocks=3)

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -835,11 +835,18 @@ class TestWorkerProcRpcRankStatus:
 
     def test_execute_rpc_returns_rank_status_envelope(self, monkeypatch):
         proc = self._make_worker_proc()
+        cpu_group = object()
 
         monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
         monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2)
+        monkeypatch.setattr(
+            diffusion_worker_module,
+            "get_world_group",
+            lambda: SimpleNamespace(cpu_group=cpu_group),
+        )
 
-        def _all_gather_object(out, local):
+        def _all_gather_object(out, local, *, group):
+            assert group is cpu_group
             out[0] = local
             out[1] = {
                 "rank": 1,

tests/engine/test_async_omni_engine_stage_init.py

@@ -2395,13 +2395,20 @@ def test_dist_stage_runtime_applies_local_dp_to_stage_config(typed):
 
 
 @pytest.mark.parametrize("num_replicas", [1, 2, 3])
-def test_typed_diffusion_replicas_share_one_config_between_planning_and_launch(mocker, num_replicas):
+@pytest.mark.parametrize("kv_owner", [None, "connector_config", "diffusion_config"])
+def test_typed_diffusion_replicas_share_one_config_between_planning_and_launch(mocker, num_replicas, kv_owner):
     from vllm_omni.config.config_factory import StageConfigFactory
     from vllm_omni.engine import stage_runtime as runtime_module
 
     stage = StageConfigFactory.create_typed_default_diffusion(
         "generic-diffusion", {"model_class_name": "QwenImagePipeline"}
     ).stage_configs[0]
+    if kv_owner is not None:
+        from vllm.config import KVTransferConfig
+
+        getattr(stage, kv_owner).kv_transfer_config = KVTransferConfig(
+            kv_connector="MooncakeConnector", kv_role="kv_consumer", engine_id="dit"
+        )
     devices = ",".join(str(i) for i in range(num_replicas))
     stage.runtime_config.devices = devices
     stage.runtime_config.num_replicas = num_replicas
@@ -2432,11 +2439,114 @@ def test_typed_diffusion_replicas_share_one_config_between_planning_and_launch(m
         for i, plan in enumerate(replicas):
             assert plan.metadata.runtime_cfg is plan.stage_cfg.runtime_config
             assert plan.stage_cfg.runtime_config.devices == str(i)
-            if num_replicas > 1:
+            if num_replicas > 1 or kv_owner is not None:
                 assert plan.stage_cfg is not stage
+            if kv_owner is not None:
+                from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
+
+                assert getattr(stage, kv_owner).kv_transfer_config.engine_id == "dit"
+                engine_args = build_engine_args_dict_from_omni_stage_config(plan.stage_cfg, "dummy-model")
+                assert engine_args["kv_transfer_config"].engine_id == f"dit-s0-r{i}"
             assert runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) is client
             assert launch.call_args.kwargs["stage_config"] is plan.stage_cfg
             assert launch.call_args.kwargs["metadata"] is plan.metadata
             assert launch.call_args.kwargs["use_inline"] is (num_replicas == 1)
 
     assert launch.call_count == 2 * num_replicas
+
+
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+@pytest.mark.parametrize("num_replicas", [1, 2])
+def test_native_kv_producer_replica_identity_and_bootstrap_are_isolated(mocker, typed, num_replicas):
+    from vllm.config import KVTransferConfig
+
+    from vllm_omni.config.omni_config import VllmOmniARStageConfig
+    from vllm_omni.config.stage_config import StagePipelineConfig
+    from vllm_omni.engine import stage_runtime as runtime_module
+
+    kv_config = KVTransferConfig(
+        kv_connector="MooncakeConnector",
+        kv_role="kv_producer",
+        engine_id="ar",
+        kv_ip="127.0.0.1",
+        kv_connector_extra_config={"bootstrap_port": 9100},
+    )
+    if typed:
+        stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="ar"))
+        stage.connector_config.kv_transfer_config = kv_config
+        runtime_cfg = stage.runtime_config
+    else:
+        runtime_cfg = OmegaConf.create({"devices": "0,1", "env": {"KEEP": "value"}})
+        stage = types.SimpleNamespace(stage_id=0, engine_args={"kv_transfer_config": kv_config}, runtime=runtime_cfg)
+        metadata = _make_llm_metadata(0)
+        mocker.patch.object(
+            runtime_module,
+            "extract_legacy_stage_metadata",
+            side_effect=lambda cfg: types.SimpleNamespace(**{**metadata.__dict__, "runtime_cfg": cfg.runtime}),
+        )
+    runtime_cfg.devices = "0,1"
+    runtime_cfg.env = {"KEEP": "value"}
+    runtime = StageRuntime(
+        stage_configs=[stage],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        async_chunk=False,
+    )
+    mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={})
+    mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None))
+    for builder in ("build_engine_args_dict", "build_engine_args_dict_from_omni_stage_config"):
+        mocker.patch.object(runtime_module, builder, return_value={})
+    mocker.patch.object(
+        runtime_module,
+        "build_vllm_config",
+        return_value=(types.SimpleNamespace(kv_transfer_config=kv_config), object),
+    )
+
+    for _ in range(2):
+        plans = runtime._build_logical_stage_init_plans(None, [num_replicas], {})
+        for i, replica in enumerate(plans[0].replicas):
+            config = replica.stage_vllm_config.kv_transfer_config
+            assert config.engine_id == f"ar-s0-r{i}"
+            assert config.kv_connector_extra_config["bootstrap_addr"] == f"http://127.0.0.1:{9100 + i}"
+            assert dict(replica.metadata.runtime_cfg.env) == {
+                "KEEP": "value",
+                "VLLM_MOONCAKE_BOOTSTRAP_PORT": str(9100 + i),
+            }
+            if typed:
+                assert replica.metadata.runtime_cfg is replica.stage_cfg.runtime_config
+        assert kv_config.engine_id == "ar"
+        assert kv_config.kv_connector_extra_config == {"bootstrap_port": 9100}
+        assert dict(runtime_cfg.env) == {"KEEP": "value"}
+
+
+@pytest.mark.parametrize(
+    "roles, sources, async_chunk, valid",
+    [
+        (["kv_producer", "kv_consumer"], [0], False, True),
+        (["kv_producer", "kv_consumer"], [0], True, False),
+        ([None, "kv_consumer"], [0], False, False),
+        (["kv_producer", None], [0], False, False),
+        ([None, "kv_producer", "kv_consumer"], [1], False, False),
+        (["kv_producer", "kv_consumer"], [1], False, False),
+        ([None, None], [0], False, True),
+    ],
+)
+def test_native_kv_topology_rejects_silent_legacy_fallback(roles, sources, async_chunk, valid):
+    plans = []
+    for stage_id, role in enumerate(roles):
+        diffusion = stage_id == len(roles) - 1
+        config = types.SimpleNamespace(kv_role=role) if role else None
+        replica = types.SimpleNamespace(
+            metadata=types.SimpleNamespace(stage_type="diffusion" if diffusion else "llm", engine_input_source=sources),
+            stage_vllm_config=None if diffusion else types.SimpleNamespace(kv_transfer_config=config),
+            stage_cfg=types.SimpleNamespace(engine_args={"kv_transfer_config": config}),
+        )
+        plans.append(types.SimpleNamespace(stage_id=stage_id, replicas=[replica]))
+    runtime = object.__new__(StageRuntime)
+    runtime._async_chunk = async_chunk
+    if valid:
+        runtime._validate_native_kv_topology(plans)
+    else:
+        with pytest.raises(ValueError, match="two-stage"):
+            runtime._validate_native_kv_topology(plans)

tests/engine/test_orchestrator.py

@@ -22,6 +22,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
 from vllm.v1.metrics.stats import IterationStats
 
 from vllm_omni.engine import OmniEngineCoreOutput
+from vllm_omni.engine.errors import NativeKVHandoffError
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
     AbortResultMessage,
@@ -1414,6 +1415,105 @@ async def test_handle_streaming_update_unknown_request_is_dropped() -> None:
     assert "req-unknown" not in orchestrator.request_states
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("native_kv", [False, True])
+async def test_add_request_attaches_native_kv_ticket_before_dispatch(mocker, native_kv) -> None:
+    sampling = SamplingParams(extra_args={"keep": "sampling"})
+    prompt = SimpleNamespace(sampling_params=SamplingParams(extra_args={"keep": "prompt"}))
+    source = StagePool(
+        0,
+        FakeStageClient(),
+        stage_vllm_config=SimpleNamespace(
+            kv_transfer_config=SimpleNamespace(kv_role="kv_producer") if native_kv else None
+        ),
+    )
+    target = StagePool(1, FakeStageClient(stage_type="diffusion", final_output=True))
+    orchestrator = Orchestrator(
+        request_async_queue=asyncio.Queue(),
+        output_async_queue=asyncio.Queue(),
+        rpc_async_queue=asyncio.Queue(),
+        stage_pools=[source, target],
+    )
+
+    async def check_dispatch(request_id, req_state, request, **kwargs):
+        assert req_state.native_kv_transfer_id == ("xfer-req" if native_kv else None)
+        for params, preserved in ((sampling, "sampling"), (request.sampling_params, "prompt")):
+            assert params.extra_args["keep"] == preserved
+            if native_kv:
+                assert params.extra_args["kv_transfer_params"] == {
+                    "transfer_id": "xfer-req",
+                    "do_remote_decode": True,
+                    "do_remote_prefill": False,
+                }
+            else:
+                assert "kv_transfer_params" not in params.extra_args
+        return 0
+
+    dispatch = mocker.patch.object(source, "submit_initial", side_effect=check_dispatch)
+    await orchestrator._handle_add_request(
+        StageSubmissionMessage(
+            type="add_request",
+            request_id="req",
+            prompt=prompt,
+            original_prompt={},
+            output_prompt_text=None,
+            sampling_params_list=[sampling, OmniDiffusionSamplingParams()],
+            final_stage_id=1,
+            preprocess_ms=0,
+            request_timestamp=0,
+            enqueue_ts=0,
+        )
+    )
+    dispatch.assert_awaited_once()
+
+
+def test_native_handoff_uses_bound_replica_and_reports_missing_binding(mocker) -> None:
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.stage_pools = [mocker.Mock()]
+    orchestrator.stage_pools[0].get_bound_client.return_value = None
+    req_state = SimpleNamespace(native_kv_transfer_id="xfer-req")
+    output = SimpleNamespace(kv_transfer_params={"num_transfer_tokens": 4})
+
+    with pytest.raises(NativeKVHandoffError, match="bound AR replica"):
+        orchestrator._diffusion_submit_kwargs("req", 0, SimpleNamespace(engine_input_source=[0]), req_state, output)
+    orchestrator.stage_pools[0].get_bound_client.assert_called_once_with("req")
+    # The pool default points at replica 0; this request actually used replica 1.
+    orchestrator.stage_pools[0].stage_vllm_config = SimpleNamespace(
+        kv_transfer_config=SimpleNamespace(engine_id="ar-0")
+    )
+    config = SimpleNamespace(engine_id="ar-1", kv_connector_extra_config={"bootstrap_addr": "http://host:8999"})
+    orchestrator.stage_pools[0].get_bound_client.return_value = SimpleNamespace(
+        vllm_config=SimpleNamespace(kv_transfer_config=config)
+    )
+    params = orchestrator._diffusion_submit_kwargs(
+        "req", 0, SimpleNamespace(engine_input_source=[0]), req_state, output
+    )
+    assert params["kv_transfer_params"]["remote_engine_id"] == "ar-1"
+    assert params["kv_transfer_params"]["remote_bootstrap_addr"] == "http://host:8999"
+    assert "remote_engine_id" not in output.kv_transfer_params
+
+
+@pytest.mark.asyncio
+async def test_native_handoff_failure_is_request_scoped(mocker):
+    orchestrator = object.__new__(Orchestrator)
+    fail = mocker.patch.object(orchestrator, "_fail_request_client_error", new_callable=mocker.AsyncMock)
+
+    def lost_binding():
+        raise NativeKVHandoffError("bound AR replica is unavailable")
+
+    assert not await orchestrator._dispatch_or_fail_request(
+        lost_binding, req_id="req", stage_id=1, operation="inter-stage forward"
+    )
+    fail.assert_awaited_once_with(
+        "req",
+        1,
+        "bound AR replica is unavailable",
+        status_code=502,
+        error_type="NativeKVHandoffError",
+        release_owners=True,
+    )
+
+
 async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> None:
     class RecordingPool:
         def __init__(self) -> None:

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -190,7 +190,7 @@ def test_legacy_default_stage_build_accepts_engine_adapter_metadata(monkeypatch)
     )[0]
     assert "api_key" not in stage_dict["engine_args"]
     stage_cfg = create_config(stage_dict)
-    metadata = SimpleNamespace(stage_id=0, cfg_kv_collect_func=None)
+    metadata = SimpleNamespace(stage_id=0, cfg_kv_collect_func=None, default_sampling_params=None)
     monkeypatch.setattr(stage_init_utils.current_omni_platform, "get_device_count", lambda: 1)
 
     config = stage_init_utils.build_diffusion_config("unused", stage_cfg, metadata)

vllm_omni/config/omni_config.py

@@ -271,6 +271,7 @@ class _ParallelEngineOverrides(_ParallelConfigEngineOverrides, total=False):
 
 class _ConnectorEngineOverrides(TypedDict, total=False):
     omni_kv_config: dict[str, Any]
+    kv_transfer_config: KVTransferConfig | dict[str, Any]
 
 
 @dataclass(frozen=True)
@@ -332,6 +333,8 @@ def _first_defined(*values: Any) -> Any:
 
 def _validate_async_chunk_support(pipeline: PipelineConfig, deploy: DeployConfig) -> None:
     has_inter_stage_edges = any(stage.input_sources for stage in pipeline.stages)
+    if deploy.async_chunk and any(stage.engine_extras.get("kv_transfer_config") for stage in deploy.stages):
+        raise ValueError("Native AR-to-DiT KV transfer requires async_chunk=False.")
     if (
         deploy.async_chunk
         and has_inter_stage_edges
@@ -605,6 +608,7 @@ class OmniStageConnectorConfig:
 
     async_chunk: bool = False
     omni_kv_config: dict[str, Any] | None = None
+    kv_transfer_config: KVTransferConfig | None = None
     stage_connector: dict[str, Any] = field(
         default_factory=lambda: {
             "name": "SharedMemoryConnector",
@@ -2078,6 +2082,7 @@ def _build_connector_config(
     return cast(Any, OmniStageConnectorConfig)(
         async_chunk=resolve_stage_async_chunk(deploy, stage_deploy),
         omni_kv_config=_copy_value(engine.get("omni_kv_config")),
+        kv_transfer_config=_copy_value(engine.get("kv_transfer_config")),
         output_connectors=_copy_value(output_connectors) if output_connectors else None,
         input_connectors=_copy_value(input_connectors) if input_connectors else None,
     )

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -23,7 +23,10 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
-from vllm_omni.core.sched.utils import omni_routed_experts_for_request
+from vllm_omni.core.sched.utils import (
+    free_kv_blocks_in_physical_order,
+    omni_routed_experts_for_request,
+)
 from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.serialization import deserialize_additional_information
 
@@ -127,6 +130,26 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Output placeholders are zero when async scheduling isn't used
         return request.num_computed_tokens - request.num_output_placeholders
 
+    def _uses_native_mooncake_connector(self) -> bool:
+        kv_config = getattr(self.vllm_config, "kv_transfer_config", None)
+        return getattr(kv_config, "kv_connector", None) == "MooncakeConnector"
+
+    def _free_request_blocks(self, request: Request) -> None:
+        """Keep native Mooncake pages coalescible without changing vLLM APIs."""
+
+        if not self._uses_native_mooncake_connector() or self.kv_cache_manager.enable_caching:
+            super()._free_request_blocks(request)
+            return
+        if not self.defer_block_free or request.last_sched_seq <= self.processed_step_seq:
+            free_kv_blocks_in_physical_order(self.kv_cache_manager, request)
+            return
+        blocks = self.kv_cache_manager.pop_blocks_for_free(request)
+        if blocks:
+            # vLLM's deferred-free drain reverses this list before returning
+            # it to BlockPool, so store the inverse of the desired order.
+            blocks.sort(key=lambda block: block.block_id, reverse=True)
+            self.deferred_frees.append((self.sched_step_seq, blocks))
+
     def _resolve_kv_connector_type(self) -> str:
         """Connector backend name for the ``kv_wait_s`` label, or ``unknown``."""
         connector_cfg = self._get_omni_kv_config_value("connector_config")
@@ -755,9 +778,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     if req_id in self.waiting_for_transfer_free:
                         req = self.requests.get(req_id)
                         if req:
-                            self.kv_cache_manager.free(req)
-                            if req_id in self.requests:
-                                del self.requests[req_id]
+                            self._free_blocks(req)
                             if req_id in self.transfer_triggered_requests:
                                 self.transfer_triggered_requests.remove(req_id)
                             self.active_kv_transfers.discard(req_id)
@@ -1041,7 +1062,36 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         getattr(self, "_inflight_prefills", set()).discard(request)
 
         # 1. Standard cleanup parts from base _free_request
-        connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
+        status = getattr(request, "status", None)
+        transfer_params = getattr(request, "kv_transfer_params", None)
+        native_transfer = (
+            transfer_params
+            and transfer_params.get("do_remote_decode")
+            and getattr(getattr(self.vllm_config, "kv_transfer_config", None), "kv_connector", None)
+            == "MooncakeConnector"
+        )
+        if native_transfer and status == RequestStatus.FINISHED_STOPPED:
+            request.status = RequestStatus.FINISHED_LENGTH_CAPPED
+        num_computed_tokens = None
+        if native_transfer:
+            # vLLM clips the block table with
+            # get_block_ids_for_computed_tokens(). Exclude Omni's optimistic
+            # async output placeholders from the physical transfer boundary.
+            num_computed_tokens = request.num_computed_tokens
+            request.num_computed_tokens = self._get_confirmed_num_computed_tokens(request)
+        try:
+            connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
+        finally:
+            if num_computed_tokens is not None:
+                request.num_computed_tokens = num_computed_tokens
+            if status is not None:
+                request.status = status
+        if native_transfer and connector_delay_free_blocks:
+            kv_xfer_params = {
+                **(kv_xfer_params or {}),
+                "transfer_id": transfer_params["transfer_id"],
+                "num_transfer_tokens": self._get_confirmed_num_computed_tokens(request),
+            }
 
         # EC Connector: mirror the KV hook (upstream v0.28 _free_request).
         # The contract requires firing before the encoder cache is freed so

vllm_omni/core/sched/utils.py

@@ -3,7 +3,9 @@
 """Shared utilities for omni schedulers."""
 
 import numpy as np
+from vllm.v1.core.kv_cache_manager import KVCacheManager
 from vllm.v1.outputs import RoutedExpertsLists
+from vllm.v1.request import Request
 
 
 def omni_routed_experts_for_request(routed_experts: RoutedExpertsLists, request) -> np.ndarray | None:
@@ -21,3 +23,19 @@ def omni_routed_experts_for_request(routed_experts: RoutedExpertsLists, request)
     mask = np.isin(routed_experts.slot_mapping, list(slot_set))
     data = routed_experts.routing_data[mask]
     return data if data.size > 0 else None
+
+
+def free_kv_blocks_in_physical_order(manager: KVCacheManager, request: Request) -> None:
+    """Return non-cached request blocks in increasing physical order.
+
+    vLLM's block pool deliberately honors the eviction-priority order passed
+    to ``free_blocks``.  For a non-caching P/D pool there is no LRU priority to
+    preserve, so physical order lets native connectors coalesce adjacent
+    source/destination pages on the next allocation.
+    """
+
+    if manager.enable_caching:
+        manager.free(request)
+        return
+    blocks = manager.pop_blocks_for_free(request)
+    manager.block_pool.free_blocks(sorted(blocks, key=lambda block: block.block_id))

vllm_omni/deploy/hunyuan-image-3-mooncakeconnector.yaml

@@ -0,0 +1,82 @@
+pipeline: hunyuan_image_3_moe
+async_chunk: false
+trust_remote_code: true
+
+stages:
+  - stage_id: 0
+    is_comprehension: true
+    # Return AR text (including CoT) alongside the final image.
+    final_output: true
+    final_output_type: text
+    num_replicas: 2
+    devices: "0,1,2,3"
+    tensor_parallel_size: 2
+    max_num_seqs: 1
+    max_num_batched_tokens: 32768
+    kv_cache_memory_bytes: 2147483648
+    gpu_memory_utilization: 0.95
+    enforce_eager: true
+    moe_backend: triton
+    env:
+      MC_TCP_ENABLE_CONNECTION_POOL: "1"
+      MC_TCP_SLICE_SIZE: "4194304"
+      VLLM_USE_FLASHINFER_SAMPLER: "0"
+    hf_overrides:
+      rope_parameters:
+        mrope_section: [0, 32, 32]
+        rope_type: default
+    kv_transfer_config:
+      kv_connector: MooncakeConnector
+      kv_role: kv_producer
+      engine_id: hunyuan-image3-ar
+      kv_ip: 127.0.0.1
+      kv_connector_extra_config:
+        mooncake_protocol: tcp
+        bootstrap_port: 8998
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      detokenize: true
+      skip_special_tokens: false
+      include_stop_str_in_output: true
+
+  - stage_id: 1
+    num_replicas: 2
+    devices: "4,5,6,7"
+    max_num_seqs: 1
+    max_num_batched_tokens: 16384
+    max_model_len: 12288
+    kv_cache_memory_bytes: 1073741824
+    enforce_eager: true
+    # Registered Mooncake KV pages must remain mapped for the connector lifetime.
+    enable_sleep_mode: false
+    moe_backend: triton
+    env:
+      MC_TCP_ENABLE_CONNECTION_POOL: "1"
+      MC_TCP_SLICE_SIZE: "4194304"
+      PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
+    distributed_executor_backend: mp
+    diffusion_kv_mode: paged_scheduler
+    diffusion_kv_max_rows_per_request: 2
+    kv_transfer_config:
+      kv_connector: MooncakeConnector
+      kv_role: kv_consumer
+      engine_id: hunyuan-image3-dit
+      kv_connector_extra_config:
+        mooncake_protocol: tcp
+        # Seconds to wait before failing a request; in-flight pages remain reserved.
+        transfer_timeout: 60.0
+    parallel_config:
+      tensor_parallel_size: 2
+      enable_expert_parallel: true
+    default_sampling_params:
+      num_inference_steps: 2
+      guidance_scale: 2.0
+      height: 512
+      width: 512
+
+edges:
+  - from: 0
+    to: 1
+    window_size: -1
+    max_inflight: 1

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -311,11 +311,17 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
         if native_impl is None:
             raise RuntimeError(f"Native attention implementation is not bound for diffusion layer {layer.layer_name!r}")
         prewrite_kv = current_omni_platform.requires_diffusion_paged_kv_prewrite()
-        read_kv_from_cache = prewrite_kv and layer.attn_backend.forward_includes_kv_cache_update
+        forward_updates_cache = layer.attn_backend.forward_includes_kv_cache_update
+        # Match vLLM's split update/attention contract. Once K/V has been
+        # written by do_kv_cache_update(), decoder attention consumes the
+        # paged cache and does not need segment-local K/V tensors. In
+        # particular, this avoids packing the same K/V projection view once
+        # for every piecewise segment.
+        read_kv_from_cache = not forward_updates_cache or prewrite_kv
         # The GPU/default path preserves native cache-update ownership. Ascend
         # prewrites through vLLM-Ascend's normal-layout writer so piecewise FIA
         # segments do not repeatedly scatter the same layer K/V.
-        if not layer.attn_backend.forward_includes_kv_cache_update or prewrite_kv:
+        if not forward_updates_cache or prewrite_kv:
             cache_update = getattr(native_impl, "do_kv_cache_update", None)
             if not callable(cache_update):
                 raise RuntimeError(
@@ -353,14 +359,11 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
             )
 
         if paged_kv_context.piecewise_plan is not None:
-            # Ascend FIA can execute identical CFG rows as one batched call
-            # per piece.  Keep that batch layout through the piecewise runner
-            # so it can concatenate row-major results instead of emitting an
-            # indexed ScatterUpdate for every segment.  The CUDA path keeps
-            # its existing output-buffer contract (including graph capture).
-            use_homogeneous_batch = (
-                current_omni_platform.is_npu() and paged_kv_context.piecewise_plan.homogeneous_batch_shape is not None
-            )
+            # Identical CFG rows can execute as one batched call per piece.
+            # Keep that batch layout through the piecewise runner so Q/K/V use
+            # strided slices instead of indexed gathers and results can be
+            # concatenated in row-major order.
+            use_homogeneous_batch = paged_kv_context.piecewise_plan.homogeneous_batch_shape is not None
             output = None
             if not use_homogeneous_batch:
                 output = torch.empty(

vllm_omni/diffusion/data.py

@@ -1225,6 +1225,8 @@ class OmniDiffusionConfig:
             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)
+            if self.enable_sleep_mode:
+                raise ValueError("Native KV transfer does not support sleep mode: registered pages must remain mapped")
 
         self.master_port = self._resolve_master_port()
         self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0)

vllm_omni/diffusion/diffusion_engine.py

@@ -603,7 +603,11 @@ class DiffusionEngine:
 
                 self._wait_for_admission_if_needed_locked()
 
-                sched_output = self.scheduler.schedule()
+                try:
+                    sched_output = self.scheduler.schedule()
+                except Exception as exc:
+                    self._fail_engine(exc)
+                    return
                 self._scheduler_num_waiting_reqs = max(int(sched_output.num_waiting_reqs), 0)
 
             self._emit_request_started_outputs(sched_output)
@@ -613,8 +617,17 @@ class DiffusionEngine:
                 continue
 
             try:
-                runner_output: BaseRunnerOutput = self.execute_fn(sched_output)  # pyright: ignore[reportAssignmentType]
+                self._prepare_kv_for_forward(sched_output)
+                runner_output: BaseRunnerOutput = (
+                    self.execute_fn(sched_output)
+                    if sched_output.scheduled_request_ids
+                    else BatchRunnerOutput.from_list([])
+                )
+                worker_execution_completed = True
             except Exception as exc:
+                if self._closed:
+                    return
+                worker_execution_completed = False
                 logger.error(
                     "Execution failed for diffusion requests %s", sched_output.scheduled_request_ids, exc_info=True
                 )
@@ -632,8 +645,23 @@ class DiffusionEngine:
 
             self._process_aborts_queue()
             self._process_rpc_queue()
-            finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
-            self._emit_outputs(finished_req_ids, sched_output.scheduled_request_ids, runner_output)
+            try:
+                finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
+                self._emit_outputs(
+                    finished_req_ids,
+                    sched_output.scheduled_request_ids,
+                    runner_output,
+                    worker_execution_completed=worker_execution_completed,
+                )
+            except Exception as exc:
+                self._fail_engine(exc)
+                return
+
+            if not sched_output.scheduled_request_ids and self.scheduler._kv_draining_requests:
+                # Only exceptional in-flight receives need idle polling. New
+                # requests / aborts wake this wait immediately.
+                with self._cv:
+                    self._cv.wait(timeout=0.01)
 
         # Engine is stopping: fail any RPCs still queued so callers don't hang.
         self._fail_pending_rpcs(RuntimeError("DiffusionEngine is shutting down."))
@@ -755,15 +783,80 @@ class DiffusionEngine:
             return
         unique_request_ids = list(dict.fromkeys(request_ids))
         if unique_request_ids:
-            self.executor.remove_diffusion_kv_requests(unique_request_ids)
+            scheduler = getattr(self, "scheduler", None)
+            target_builder = getattr(scheduler, "get_diffusion_kv_cleanup_targets", None)
+            cleanup_targets = target_builder(unique_request_ids) if target_builder is not None else unique_request_ids
+            try:
+                if cleanup_targets:
+                    self.executor.remove_diffusion_kv_requests(cleanup_targets)
+            except Exception as exc:
+                self._fail_engine(exc)
+                raise
+
+    def _prepare_kv_for_forward(self, sched_output: DiffusionSchedulerOutput) -> None:
+        if getattr(sched_output, "kv_connector_metadata", None) is None:
+            return
+        try:
+            output = self.executor.prepare_kv_for_forward(sched_output)
+            assert output is not None
+            self.scheduler.update_kv_connector_output(output)
+            if not self.abort_queue.empty():
+                self._process_aborts_queue()
+            incomplete = sched_output.kv_transfer_request_ids - (output.finished_recving or set())
+            sched_output.finished_req_ids.update(self.scheduler.fail_incomplete_kv_loads(incomplete))
+            # Timed-out/cancelled requests must never reach model execution.
+            terminal = {
+                rid
+                for rid in sched_output.scheduled_request_ids
+                if (state := self.scheduler.get_request_state(rid)) is not None and state.is_finished()
+            }
+            if terminal:
+                sched_output.finished_req_ids.update(terminal)
+                sched_output.scheduled_new_reqs = [
+                    req for req in sched_output.scheduled_new_reqs if req.request_id not in terminal
+                ]
+                sched_output.scheduled_cached_reqs.request_ids = [
+                    rid for rid in sched_output.scheduled_cached_reqs.request_ids if rid not in terminal
+                ]
+                # is_empty / started-output handling may have already cached it.
+                sched_output.__dict__.pop("scheduled_request_ids", None)
+            drained = self.scheduler.completed_kv_drains()
+            if drained:
+                self._remove_diffusion_kv_requests(drained)
+                self.scheduler.release_kv_drains(drained)
+        except Exception as exc:
+            self._fail_engine(exc)
+            raise
+
+    def _fail_engine(self, exc: Exception) -> None:
+        if getattr(self, "_shutdown_complete", False):
+            return
+        logger.error("Diffusion engine failed; stopping workers before releasing KV pages", exc_info=exc)
+        with self._cv:
+            self._closed = True
+            if self.stop_event is not None:
+                self.stop_event.set()
+            streams = list(self._out_streams.values())
+            self._cv.notify_all()
+        for stream in streams:
+            self._put_queue_output(stream, DiffusionOutput.from_exception(exc))
+        self._fail_pending_rpcs(exc)
+        try:
+            self.executor.shutdown()
+        finally:
+            try:
+                self.scheduler.close()
+            finally:
+                self._shutdown_complete = True
 
     def _emit_finished_outputs(
         self,
         finished_ids: set[str],
         runner_output: BaseRunnerOutput | None = None,
         missing_result_error: str = "Diffusion execution finished without a final output",
+        worker_cleaned_ids: set[str] | None = None,
     ) -> None:
-        self._remove_diffusion_kv_requests(finished_ids)
+        self._remove_diffusion_kv_requests(finished_ids - (worker_cleaned_ids or set()))
         for rid in finished_ids:
             if runner_output is not None:
                 _output = runner_output.get_request_output(rid)
@@ -777,13 +870,21 @@ class DiffusionEngine:
         finished_ids: set[str],
         scheduled_request_ids: list[str],
         runner_output: BaseRunnerOutput,
+        *,
+        worker_execution_completed: bool = False,
     ) -> None:
         """Emit output chunks for every request through the unified output stream."""
+        finished_ids = set(finished_ids)
+        worker_cleaned_ids = finished_ids.intersection(scheduled_request_ids) if worker_execution_completed else set()
         if self.execution_mode != DiffusionExecutionMode.STEP_BATCH:
-            self._emit_finished_outputs(finished_ids, runner_output)
+            self._emit_finished_outputs(
+                finished_ids,
+                runner_output,
+                worker_cleaned_ids=worker_cleaned_ids,
+            )
             return
 
-        self._remove_diffusion_kv_requests(finished_ids)
+        self._remove_diffusion_kv_requests(finished_ids - worker_cleaned_ids)
 
         delivered_finished_req_ids: set[str] = set()
 
@@ -975,7 +1076,11 @@ class DiffusionEngine:
             # keep scheduling and executing until the target request is finished
             while True:
                 self._process_aborts_queue()
-                sched_output = self.scheduler.schedule()
+                try:
+                    sched_output = self.scheduler.schedule()
+                except Exception as exc:
+                    self._fail_engine(exc)
+                    raise
                 if sched_output.is_empty:
                     if target_request_id in sched_output.finished_req_ids:
                         self._remove_diffusion_kv_requests([target_request_id])
@@ -986,12 +1091,19 @@ class DiffusionEngine:
 
                 # NOTE: add_req_and_wait_for_response() is synchronous, will be only called
                 # within _dummy_run, only one request will be scheduled
-                request_id = sched_output.scheduled_request_ids[0]
+                request_id = target_request_id
                 try:
-                    runner_output: BaseRunnerOutput = self.execute_fn(sched_output)  # pyright: ignore[reportAssignmentType]
+                    self._prepare_kv_for_forward(sched_output)
+                    runner_output: BaseRunnerOutput = (
+                        self.execute_fn(sched_output)
+                        if sched_output.scheduled_request_ids
+                        else BatchRunnerOutput.from_list([])
+                    )
                 except EngineDeadError:
                     raise
                 except Exception as exc:
+                    if self._closed:
+                        raise
                     logger.error("Execution failed for diffusion request %s", request_id, exc_info=True)
                     runner_output = RunnerOutput(
                         request_id=request_id,
@@ -1005,7 +1117,11 @@ class DiffusionEngine:
                 finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
 
                 # sync func should receive one result
-                if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1:
+                if (
+                    sched_output.scheduled_request_ids
+                    and not isinstance(runner_output, RunnerOutput)
+                    and len(runner_output) != 1
+                ):
                     raise ValueError("Sync func should receive one result at one time")
                 if target_request_id in finished_req_ids:
                     self._remove_diffusion_kv_requests([target_request_id])
@@ -1100,7 +1216,8 @@ class DiffusionEngine:
         executes one request from the collective wave.
 
         Hunyuan is currently the only model integrated with
-        ``paged_scheduler``; 1024x1024, enabled CFG, and the maximum advertised
+        ``paged_scheduler``. The stage's configured default resolution (or
+        1024x1024 when absent), enabled CFG, and the maximum advertised
         reference-image count exercise its first-step activation peak.
         Admission compares each preprocessed request's CFG count and tokenized
         sequence/target shape with the resulting per-request profile envelope.
@@ -1117,9 +1234,12 @@ class DiffusionEngine:
         model_class_name = self.od_config.model_class_name
         if model_class_name is None:
             raise RuntimeError("Diffusion KV profiling requires a resolved model_class_name")
+        additional_config = getattr(self.od_config, "additional_config", {}) or {}
+        profile_height = additional_config.get("diffusion_kv_profile_height", 1024)
+        profile_width = additional_config.get("diffusion_kv_profile_width", 1024)
         request = self._make_dummy_request(
-            height=1024,
-            width=1024,
+            height=profile_height,
+            width=profile_width,
             guidance_scale=5.0,
             num_image_inputs=get_dummy_run_num_image_inputs(model_class_name),
         )
@@ -1335,7 +1455,11 @@ class DiffusionEngine:
 
     def _process_aborts_queue(self) -> None:
         with self._cv:
-            self._drain_abort_queue()
+            try:
+                self._drain_abort_queue()
+            except Exception as exc:
+                self._fail_engine(exc)
+                raise
 
     def _drain_abort_queue(self) -> None:
         if self.abort_queue.empty():
@@ -1352,11 +1476,10 @@ class DiffusionEngine:
         request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids)
         request_ids = list(dict.fromkeys(request_ids))
 
-        self._remove_diffusion_kv_requests(request_ids)
-
         for request_id in request_ids:
             if self.scheduler.get_request_state(request_id) is not None:
                 self.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
+        self._remove_diffusion_kv_requests(request_ids)
 
     def _finalize_finished_request(
         self,

vllm_omni/diffusion/diffusion_kv/kv_connector.py

@@ -1,35 +1,40 @@
 # 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.
-"""
+"""Assemble native vLLM KV connectors for Scheduler-owned diffusion pages."""
 
 from __future__ import annotations
 
+import math
+import time
+from collections import Counter
 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.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.outputs import KVConnectorOutput
 
 if TYPE_CHECKING:
     from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+    from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
 
     from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+    from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
 
 logger = init_logger(__name__)
 
 
+class KVTransferRegistrationError(ValueError):
+    """A failed handoff whose destination pages have not been dispatched."""
+
+
 def mint_transfer_id(request_id: str) -> str:
     """Return the stable ticket shared by one AR -> Diffusion handoff."""
 
@@ -110,6 +115,8 @@ def parse_kv_transfer_config(value: object | None) -> KVTransferConfig | None:
 
 def create_scheduler_kv_connector(
     od_config: OmniDiffusionConfig,
+    kv_cache_config: KVCacheConfig | None = None,
+    vllm_config: VllmConfig | None = None,
 ) -> KVConnectorBase_V1 | None:
     """Create a Scheduler-role connector when native config is present."""
 
@@ -121,14 +128,15 @@ def create_scheduler_kv_connector(
             f"Diffusion native kv_transfer_config must be KVTransferConfig, got {type(kv_transfer_config)!r}"
         )
 
-    vllm_config = _build_native_vllm_config(od_config)
+    if kv_cache_config is None or vllm_config is None:
+        raise ValueError("KV transfer requires paged_scheduler and a native cache configuration")
     connector = KVConnectorFactory.create_connector(
         config=vllm_config,
         role=KVConnectorRole.SCHEDULER,
-        kv_cache_config=_empty_kv_cache_config(),
+        kv_cache_config=kv_cache_config,
     )
     logger.info(
-        "Created KV connector stub (SCHEDULER role): connector=%s engine_id=%s",
+        "Created KV connector (SCHEDULER role): connector=%s engine_id=%s",
         kv_transfer_config.kv_connector,
         kv_transfer_config.engine_id,
     )
@@ -140,9 +148,27 @@ def init_worker_kv_connector(vllm_config: VllmConfig, kv_cache_config: KVCacheCo
 
     if vllm_config.kv_transfer_config is None:
         return
-    ensure_kv_transfer_initialized(vllm_config, kv_cache_config)
+    import vllm.distributed.parallel_state as parallel_state
+
+    from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
+
+    # vLLM 0.29 Mooncake snapshots tp_rank/tp_size in its Worker constructor;
+    # consumer receive/region mapping use those saved values, not get_tp_group.
+    # Restrict this compatibility bridge to that connector and restore the
+    # model's TP group even when construction fails. Never remap model TP for
+    # the connector's entire lifetime (it would change model collectives).
+    tp_group = parallel_state._TP
+    parallel_config = getattr(vllm_config, "parallel_config", None)
+    if getattr(parallel_config, "prefill_context_parallel_size", 1) > 1 and tp_group.world_size == 1:
+        if vllm_config.kv_transfer_config.kv_connector != "MooncakeConnector":
+            raise ValueError("Native SP KV transfer currently requires MooncakeConnector")
+        parallel_state._TP = get_sp_group()
+    try:
+        ensure_kv_transfer_initialized(vllm_config, kv_cache_config)
+    finally:
+        parallel_state._TP = tp_group
     logger.info(
-        "Initialized KV connector stub (WORKER role): connector=%s engine_id=%s",
+        "Initialized KV connector (WORKER role): connector=%s engine_id=%s",
         vllm_config.kv_transfer_config.kv_connector,
         vllm_config.kv_transfer_config.engine_id,
     )
@@ -156,26 +182,165 @@ def shutdown_kv_connector(*, scheduler_connector: KVConnectorBase_V1 | None = No
     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 prepare_kv_requests(requests: tuple[DiffusionKVRequest, ...], params: Mapping[str, Any]) -> None:
+    for request in requests:
+        request.kv_transfer_params = dict(params)
+        num_tokens = min(len(request.prompt_token_ids), params["num_transfer_tokens"])
+        request.prompt_token_ids = request.prompt_token_ids[:num_tokens]
+        request.num_prompt_tokens = num_tokens
 
 
-def _build_native_vllm_config(od_config: OmniDiffusionConfig) -> VllmConfig:
-    from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
+def install_mooncake_cfg_fanout(connector: KVConnectorBase_V1) -> None:
+    """Account for multiple CFG destinations on one native Mooncake ticket.
 
-    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
+    vLLM 0.29 initializes ``need_send`` to the paired consumer rank count,
+    but increments ``sent`` per consumer request. Adapt only this producer
+    instance, after KV initialization and before serving requests. Keep the
+    upstream address planning, transport, timeout and completion code intact.
+    """
+    from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import MooncakeConnector
+
+    if not isinstance(connector, MooncakeConnector):
+        return
+    worker = connector.connector_worker
+    if worker is None or worker.is_kv_consumer or getattr(worker, "_omni_cfg_fanout_installed", False):
+        return
+    build_transfer_params = worker._build_transfer_params
+
+    async def build_with_fanout(ready_reqs, agent_meta, local_regions, remote_regions):
+        # commit_kv_load registers all rows before build_connector_meta. Thus
+        # each rank's message contains the complete row set for each ticket,
+        # including empty-block notifications. Count the full message, not
+        # just the subset whose producer-ready events fired in this iteration.
+        rows_per_ticket = Counter(transfer_id for transfer_id, _ in agent_meta.req_blocks.values())
+        tickets = {send_meta.transfer_id: send_meta for _, send_meta in ready_reqs}
+        for transfer_id, send_meta in tickets.items():
+            rows = rows_per_ticket[transfer_id]
+            previous = getattr(send_meta, "_omni_cfg_rows", None)
+            if rows < 1 or (previous is not None and previous != rows):
+                # Use the upstream transfer-error path: it decrements sending
+                # without incrementing sent, so a malformed rank cannot free
+                # source pages still needed by another rank.
+                return [], [], [], [req_id for req_id, _ in ready_reqs], "Inconsistent Mooncake CFG receive count"
+        # No await between checking and updating: all handlers run on the
+        # sender loop. The marker belongs to the ticket, not to a request-global
+        # variable; it disappears with the upstream ticket on completion/expiry.
+        for transfer_id, send_meta in tickets.items():
+            if getattr(send_meta, "_omni_cfg_rows", None) is None:
+                send_meta.need_send *= rows_per_ticket[transfer_id]
+                send_meta._omni_cfg_rows = rows_per_ticket[transfer_id]
+        return await build_transfer_params(ready_reqs, agent_meta, local_regions, remote_regions)
+
+    worker._build_transfer_params = build_with_fanout  # type: ignore[method-assign]
+    setattr(worker, "_omni_cfg_fanout_installed", True)
+
+
+def commit_kv_load(
+    connector: KVConnectorBase_V1,
+    manager: KVCacheManager,
+    requests: tuple[DiffusionKVRequest, ...],
+    matched_tokens: list[int],
+) -> set[str]:
+    expected = set()
+    allocations = []
+    # All rows sharing a transfer_id must reach the same connector metadata:
+    # the producer uses this complete per-rank row set for fan-out accounting.
+    # Validate every CFG row before mutating any connector state.
+    for request, num_tokens in zip(requests, matched_tokens, strict=True):
+        blocks = manager.get_blocks(request.request_id)
+        # Mooncake's producer advertises complete physical blocks. Keep every
+        # CFG receiver on that same block boundary so the native connector can
+        # use its normal one-to-one page mapping even when the branches have
+        # different logical prefix lengths. Bytes beyond ``num_tokens`` are
+        # not marked computed and are overwritten by the DiT prefill.
+        transfer_tokens = num_tokens
+        if num_tokens > 0 and request.kv_transfer_params is not None:
+            transfer_tokens = request.kv_transfer_params["num_transfer_tokens"]
+            if type(transfer_tokens) is not int or not num_tokens <= transfer_tokens <= request.num_tokens:
+                raise KVTransferRegistrationError(
+                    "Diffusion KV transfer boundary must cover the reusable prefix "
+                    f"without exceeding the allocated sequence: reusable={num_tokens}, "
+                    f"transfer={transfer_tokens!r}, allocated={request.num_tokens}"
+                )
+        prefix_blocks = KVCacheBlocks(
+            tuple(
+                group[: (transfer_tokens + spec.kv_cache_spec.block_size - 1) // spec.kv_cache_spec.block_size]
+                for group, spec in zip(blocks.blocks, manager.kv_cache_config.kv_cache_groups, strict=True)
+            )
+        )
+        allocations.append((request, prefix_blocks, num_tokens))
+    try:
+        for request, prefix_blocks, num_tokens in allocations:
+            connector.update_state_after_alloc(request, prefix_blocks, num_tokens)
+            request.num_computed_tokens = num_tokens
+            if request.kv_transfer_params is not None:
+                expected.add(request.request_id)
+    except Exception as exc:
+        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import MooncakeConnector
+        from vllm.v1.request import RequestStatus
+
+        if not isinstance(connector, MooncakeConnector):
+            # Unknown connectors may have started I/O during registration.
+            raise
+        # Mooncake only queues metadata here; no Worker has seen addresses yet.
+        # Re-arm its pre-scheduling abort hook to replace even partially
+        # registered CFG receives with empty-block notifications to the producer.
+        for request in requests:
+            if request.kv_transfer_params is not None:
+                request.kv_transfer_params["do_remote_prefill"] = True
+                request.status = RequestStatus.FINISHED_ABORTED
+                delay_free, _ = connector.request_finished(request, [])
+                if delay_free:
+                    raise RuntimeError("Connector could not cancel an undispatched KV load") from exc
+        raise KVTransferRegistrationError(f"Could not register diffusion KV receive: {exc}") from exc
+    return expected
+
+
+def wait_for_kv_load(
+    active_connector: ActiveKVConnector, scheduler_output: DiffusionSchedulerOutput, timeout: float
+) -> KVConnectorOutput:
+    connector = active_connector.kv_connector
+    active_connector.pre_forward(scheduler_output)
+    pending = set(scheduler_output.kv_transfer_request_ids)
+    finished_ids = scheduler_output.kv_finished_request_ids
+    if not pending:
+        output = active_connector.post_forward(finished_ids)
+        if output.invalid_block_ids:
+            raise RuntimeError("Diffusion KV connector reported invalid remote pages")
+        return output
+
+    received, sent = set(), set()
+    deadline = time.monotonic() + timeout
+    while True:
+        finished_sending, finished_recving = connector.get_finished(finished_ids)
+        sent.update(finished_sending or ())
+        received.update(finished_recving or ())
+        pending.difference_update(received)
+        if connector.get_block_ids_with_load_errors():
+            raise RuntimeError("Diffusion KV connector reported invalid remote pages")
+        if not pending:
+            break
+        if time.monotonic() >= deadline:
+            # A timeout is not a cancellation acknowledgement. Return partial
+            # completion; Scheduler quarantines the remaining destination pages.
+            break
+        time.sleep(0.001)
+    output = active_connector.post_forward(finished_ids)
+    output.finished_sending = sent | (output.finished_sending or set())
+    output.finished_recving = received | (output.finished_recving or set())
+    if output.invalid_block_ids:
+        raise RuntimeError("Diffusion KV connector reported invalid remote pages")
+    return output
 
 
 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")
+    extra_config = config.kv_connector_extra_config or {}
+    timeout = extra_config.get("transfer_timeout", 60.0)
+    if isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or not math.isfinite(timeout) or timeout <= 0:
+        raise ValueError("Diffusion transfer_timeout must be a positive finite number of seconds")
     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"):

vllm_omni/diffusion/diffusion_kv/manager.py

@@ -8,6 +8,7 @@ from collections.abc import Sequence
 from vllm.v1.core.kv_cache_manager import KVCacheManager
 from vllm.v1.kv_cache_interface import KVCacheConfig
 
+from vllm_omni.core.sched.utils import free_kv_blocks_in_physical_order
 from vllm_omni.diffusion.diffusion_kv.metadata import (
     DiffusionKVMetadata,
     DiffusionKVSequenceMetadata,
@@ -179,7 +180,7 @@ class DiffusionKVCacheManager:
         requests = self._requests.pop(public_request_id, ())
         self._metadata.pop(public_request_id, None)
         for request in reversed(requests):
-            self.native_manager.free(request)
+            free_kv_blocks_in_physical_order(self.native_manager, request)
             self._internal_request_ids.discard(request.request_id)
 
     def close(self) -> None:
@@ -188,4 +189,4 @@ class DiffusionKVCacheManager:
 
     def _rollback(self, requests: Sequence[DiffusionKVRequest]) -> None:
         for request in reversed(requests):
-            self.native_manager.free(request)
+            free_kv_blocks_in_physical_order(self.native_manager, request)

vllm_omni/diffusion/diffusion_kv/metadata.py

@@ -41,6 +41,7 @@ class DiffusionKVSequenceMetadata:
     seq_len: int
     block_ids: tuple[list[int], ...]
     context_ids: tuple[str, ...] = ()
+    num_computed_tokens: int = 0
 
 
 @dataclass

vllm_omni/diffusion/diffusion_kv/model_runner_backend.py

@@ -190,6 +190,10 @@ class DiffusionKVModelRunnerBackend:
             )
         return ulysses_degree
 
+    @property
+    def kv_caches_by_layer(self) -> dict[str, torch.Tensor]:
+        return {name: layer.kv_cache for name, layer in self._kv_cache_layer_adapters.items()}
+
     def _get_max_rows_per_request(self) -> int:
         max_rows = getattr(self.od_config, "diffusion_kv_max_rows_per_request", None)
         if type(max_rows) is not int or max_rows <= 0:
@@ -245,7 +249,9 @@ class DiffusionKVModelRunnerBackend:
             for capacity, block_size in zip(unaligned_max_num_blocks_per_group, block_sizes, strict=True)
         ]
 
-        max_num_reqs = max_num_seqs * max_rows_per_request
+        max_num_reqs = max_num_seqs
+        if max_num_reqs % max_rows_per_request:
+            raise ValueError("scheduler_config.max_num_seqs must be a multiple of diffusion_kv_max_rows_per_request")
         max_num_batched_tokens = getattr(scheduler_config, "max_num_batched_tokens", None)
         if type(max_num_batched_tokens) is not int or max_num_batched_tokens <= 0:
             raise ValueError("scheduler_config.max_num_batched_tokens must be a positive integer")
@@ -254,9 +260,6 @@ class DiffusionKVModelRunnerBackend:
         cp_rank = get_dcp_group().rank_in_group if cp_size > 1 else 0
         cp_interleave = parallel_config.cp_kv_cache_interleave_size
 
-        # Native metadata builders size per-row buffers from max_num_seqs.
-        # A diffusion public request can occupy several sequence/context rows.
-        scheduler_config.max_num_seqs = max_num_reqs
         kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
         previous_adapter_caches = {
             layer_name: adapter.kv_cache for layer_name, adapter in self._kv_cache_layer_adapters.items()
@@ -303,7 +306,6 @@ class DiffusionKVModelRunnerBackend:
                     self.vllm_config,
                 )
         except Exception:
-            scheduler_config.max_num_seqs = max_num_seqs
             for layer_name, previous_cache in previous_adapter_caches.items():
                 self._kv_cache_layer_adapters[layer_name].kv_cache = previous_cache
             raise
@@ -591,7 +593,7 @@ class DiffusionKVModelRunnerBackend:
                     )
         raise RuntimeError(f"Diffusion KV request state is missing logical length for {identity!r}")
 
-    def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int:
+    def remove_diffusion_kv_requests(self, request_ids: list[str | tuple[str, int]]) -> int:
         """Retire Worker rows without logically freeing Scheduler-owned blocks."""
         if (
             not is_scheduler_paged_kv_mode(
@@ -600,11 +602,25 @@ class DiffusionKVModelRunnerBackend:
             or self.block_tables is None
         ):
             return 0
-        request_id_set = set(request_ids)
+        request_generations = {
+            request_id: generation
+            for item in request_ids
+            if isinstance(item, tuple)
+            for request_id, generation in (item,)
+            if isinstance(request_id, str) and type(generation) is int
+        }
+        request_id_set = {item if isinstance(item, str) else item[0] for item in request_ids}
         identities_and_rows = [
             (identity, row)
             for identity, row in self._diffusion_kv_identity_to_row.items()
             if identity[0] in request_id_set
+            and (
+                identity[0] not in request_generations
+                or (
+                    (request_state := self._diffusion_kv_request_states.get(identity[0])) is not None
+                    and request_state.generation == request_generations[identity[0]]
+                )
+            )
         ]
         rows = [row for _, row in identities_and_rows]
         if rows:
@@ -620,9 +636,10 @@ class DiffusionKVModelRunnerBackend:
             ]
             self._apply_rows(rows, installs)
 
+        removed_request_ids = {identity[0] for identity, _ in identities_and_rows}
         for identity, _ in identities_and_rows:
             del self._diffusion_kv_identity_to_row[identity]
-        for request_id in request_id_set:
+        for request_id in removed_request_ids:
             self._diffusion_kv_request_states.pop(request_id, None)
         self._diffusion_kv_free_rows.extend(rows)
         self._diffusion_kv_free_rows.sort(reverse=True)

vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py

@@ -84,12 +84,14 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
         attention_config = vllm_config.attention_config
         previous_backend = attention_config.backend
         previous_backend_per_kind = attention_config.backend_per_kind
+        previous_pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
         try:
             # This is a portable vLLM backend request. The active platform
             # resolves it to its native implementation, such as FlashAttention
             # on CUDA or AscendAttentionBackend on NPU.
             attention_config.backend = AttentionBackendEnum.FLASH_ATTN
             attention_config.backend_per_kind = {}
+            vllm_config.parallel_config.prefill_context_parallel_size = 1
             with set_current_vllm_config(vllm_config):
                 attn_backend = get_attn_backend(
                     head_size=spec.head_size,
@@ -101,6 +103,7 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
         finally:
             attention_config.backend = previous_backend
             attention_config.backend_per_kind = previous_backend_per_kind
+            vllm_config.parallel_config.prefill_context_parallel_size = previous_pcp_size
         attn_backend = current_omni_platform.get_diffusion_paged_kv_attn_backend(
             attn_backend,
             ulysses_degree=ulysses_degree,
@@ -758,6 +761,19 @@ class DiffusionPagedAttentionAdapter:
                 slot_mappings=slot_mappings,
                 causal=(row_segments[0].mode == "causal"),
             )
+            # FA2's ``num_splits=0`` auto policy selects SplitKV for these
+            # already-small piecewise calls. Reuse vLLM's supported metadata
+            # control to select the single-pass kernel and avoid its extra
+            # partial-output/reduction work. Leave newer native backends and
+            # metadata without this control untouched.
+            for layer_name, native_metadata in segment_metadata.items():
+                layer = self.layers.get(layer_name)
+                if (
+                    layer is not None
+                    and getattr(layer.impl, "vllm_flash_attn_version", None) == 2
+                    and hasattr(native_metadata, "max_num_splits")
+                ):
+                    native_metadata.max_num_splits = 1
             # FA3's full-CUDA-graph metadata builder reuses one persistent
             # scheduler buffer across builds. Piecewise attention prepares all
             # segments before executing any of them, so each segment needs its
@@ -901,9 +917,17 @@ class DiffusionPagedAttentionAdapter:
                 raise KeyError(
                     f"No piecewise native attention metadata was built for diffusion layer {layer_name!r}"
                 ) from exc
+        # Identical rows are sliced and packed once per segment by the
+        # homogeneous piecewise runner.  Keeping the projection view here
+        # avoids first copying the complete Q tensor only to copy each segment
+        # again immediately afterwards.  Other native paths still require one
+        # packed query buffer.
+        query_input = query_flat
+        if piecewise_plan is None or piecewise_plan.homogeneous_batch_shape is None:
+            query_input = query_flat.contiguous()
         return DiffusionPagedAttentionContext(
             layer=layer,
-            query=query_flat.contiguous(),
+            query=query_input,
             key_write=key_flat,
             value_write=value_flat,
             slot_mapping=slot_mapping,

vllm_omni/diffusion/diffusion_kv/request.py

@@ -113,8 +113,6 @@ class DiffusionKVRequest:
         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

vllm_omni/diffusion/executor/abstract.py

@@ -1,13 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable
+from dataclasses import replace
 from typing import TYPE_CHECKING, Any
 
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
+from vllm.v1.outputs import KVConnectorOutput
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.sched.interface import CachedRequestData
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -159,7 +165,7 @@ class DiffusionExecutor(ABC):
         # ranks are not silently dropped.
         self.collective_rpc("set_kv_cache_configs", args=(kv_cache_configs, resolved_max_model_len))
 
-    def remove_diffusion_kv_requests(self, request_ids: list[str]) -> None:
+    def remove_diffusion_kv_requests(self, request_ids: list[str | tuple[str, int]]) -> None:
         """Clear request rows on every Worker after Scheduler retirement."""
 
         unique_request_ids = list(dict.fromkeys(request_ids))
@@ -170,6 +176,35 @@ class DiffusionExecutor(ABC):
             args=(unique_request_ids,),
         )
 
+    def prepare_kv_for_forward(self, scheduler_output: DiffusionSchedulerOutput) -> KVConnectorOutput | None:
+        if scheduler_output.kv_connector_metadata is None:
+            return
+        transfer_output = replace(
+            scheduler_output,
+            scheduled_new_reqs=[],
+            scheduled_cached_reqs=CachedRequestData.make_empty(),
+            kv_prefetch_job=None,
+        )
+        outputs: list[KVConnectorOutput] = self.collective_rpc(
+            "prepare_kv_for_forward",
+            args=(transfer_output,),
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        if len(outputs) != self.od_config.num_gpus or any(output.invalid_block_ids for output in outputs):
+            # Missing ranks / invalid pages cannot establish safe ownership.
+            raise RuntimeError("Diffusion KV receive failed on one or more ranks")
+        completed = getattr(self, "_kv_receive_completed_ranks", {})
+        for rank, output in enumerate(outputs):
+            for request_id in output.finished_recving or ():
+                completed.setdefault(request_id, set()).add(rank)
+        finished = {request_id for request_id, ranks in completed.items() if len(ranks) == len(outputs)}
+        for request_id in finished:
+            del completed[request_id]
+        self._kv_receive_completed_ranks = completed
+        outputs[0].finished_recving = finished
+        return outputs[0]
+
     @abstractmethod
     def shutdown(self) -> None:
         """Shutdown the executor and release resources."""

vllm_omni/diffusion/forward_context.py

@@ -282,6 +282,47 @@ def set_forward_context_denoise_step_idx(step_idx: int | None) -> None:
                 ensure_active(step_idx)
 
 
+def get_paged_kv_computed_tokens() -> tuple[int, ...]:
+    runtime = _forward_context.paged_kv_runtime if _forward_context is not None else None
+    if runtime is None:
+        return ()
+    return tuple(row.kv_start_pos for row in runtime.metadata.prefill_rows)
+
+
+@contextmanager
+def paged_kv_prefill(sequence_id: int, num_tokens: int):
+    from dataclasses import replace
+
+    runtime = get_forward_context().paged_kv_runtime
+    if runtime is None:
+        raise RuntimeError("Paged KV prefill requires an active runtime")
+    rows = runtime.metadata.prefill_rows
+    matching_rows = [row for row in rows if row.sequence_id == sequence_id]
+    if len(matching_rows) != 1:
+        raise ValueError(
+            f"Paged KV prefill requires exactly one active row for sequence {sequence_id}; found {len(matching_rows)}"
+        )
+    row = matching_rows[0]
+    if type(num_tokens) is not int or not row.kv_start_pos < num_tokens <= row.seq_len:
+        raise ValueError(
+            "Paged KV prefill target must extend the active prefix without exceeding its allocation: "
+            f"start={row.kv_start_pos}, target={num_tokens!r}, allocated={row.seq_len}"
+        )
+    prefill = replace(row, query_len=num_tokens - row.kv_start_pos, seq_len=num_tokens)
+    batch = runtime.adapter.prepare_batch((prefill,))
+    with runtime.adapter.activate(batch):
+        yield
+    runtime.metadata = replace(
+        runtime.metadata,
+        prefill_rows=tuple(
+            replace(item, kv_start_pos=num_tokens, query_len=item.seq_len - num_tokens)
+            if item.sequence_id == sequence_id
+            else item
+            for item in rows
+        ),
+    )
+
+
 def set_forward_context_denoise_timestep(timestep: float | None) -> None:
     """Set the normalized (descending, 1 -> 0) denoise timestep.
 

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -75,7 +75,11 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelOutput,
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
+from vllm_omni.diffusion.forward_context import (
+    get_paged_kv_computed_tokens,
+    paged_kv_prefill,
+    set_forward_context_denoise_step_idx,
+)
 from vllm_omni.diffusion.layers.fused_moe import FusedMoE
 from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
@@ -1097,8 +1101,6 @@ class ImageKVCacheManager(nn.Module):
         layer resolves the actual CUDA or Ascend paged kernel.
         """
 
-        if uncond_cfg_prefill:
-            raise RuntimeError("Hunyuan negative-CFG prefill must run before paged row activation")
         if self._injected_ar_kv is not None:
             raise NotImplementedError("Hunyuan Scheduler-paged KV does not support imported AR KV")
         if not query_lens or len(seq_lens) != len(query_lens):
@@ -1109,7 +1111,6 @@ class ImageKVCacheManager(nn.Module):
         self.clear_legacy_prompt_kv_cache()
         bs = len(query_lens)
         q_len = query_lens[0]
-        seq_len = seq_lens[0]
         assert query.shape[0] == bs * q_len, f"{query.shape[0]} != {bs * q_len}"
 
         head_num_per_rank = query.shape[1]
@@ -1121,22 +1122,19 @@ class ImageKVCacheManager(nn.Module):
         value = value.reshape(bs, q_len, kv_head_num_per_rank, head_dim)
 
         joint_text_query = joint_text_key = joint_text_value = None
-        if self.sp_size > 1 and first_step:
+        if self.sp_size > 1 and uncond_cfg_prefill:
+            joint_text_query, joint_text_key, joint_text_value = query, key, value
+            query, key, value = query[:, :0], key[:, :0], value[:, :0]
+        elif self.sp_size > 1 and first_step:
             if shard_image_size is None or shard_image_size <= 0:
                 raise ValueError("Hunyuan paged Ulysses requires a positive local image shard size")
-            local_prompt_len = seq_len - shard_image_size
             joint_query_len = query.shape[1] - shard_image_size
-            if local_prompt_len != joint_query_len:
-                raise ValueError(
-                    "Hunyuan paged Ulysses prompt layout mismatch: "
-                    f"key_prompt={local_prompt_len}, query_prompt={joint_query_len}"
-                )
             joint_text_query = query[:, :joint_query_len]
-            joint_text_key = key[:, :local_prompt_len]
-            joint_text_value = value[:, :local_prompt_len]
+            joint_text_key = key[:, :joint_query_len]
+            joint_text_value = value[:, :joint_query_len]
             query = query[:, joint_query_len:]
-            key = key[:, local_prompt_len:]
-            value = value[:, local_prompt_len:]
+            key = key[:, joint_query_len:]
+            value = value[:, joint_query_len:]
 
         if joint_text_query is None:
             attn_metadata = AttentionMetadata(full_attn_spans=full_attn_spans)
@@ -1254,6 +1252,8 @@ class ImageKVCacheManager(nn.Module):
     ) -> torch.Tensor:
         self.image_token_len = kwargs.get("num_image_tokens")
         if self.attn.is_paged_kv_active():
+            # A native connector fills the same Scheduler-owned pages as
+            # local prefill. Imported prefixes need no dense reconstruction.
             full_attn_spans = kwargs.get("full_attn_spans")
             if full_attn_spans is None:
                 raise ValueError("Hunyuan Scheduler-paged KV requires full_attn_spans metadata")
@@ -1837,7 +1837,9 @@ class HunYuanAttention(nn.Module):
             num_kv_heads=self.num_kv_heads,
             scaling=self.scaling,
             image_token_len=4097,
-            prefix=f"{prefix}.image_attn",
+            # Native Mooncake pairs this canonical KV layer name with the AR
+            # HunYuanAttention: model.layers.N.self_attn.attn. Keep both in sync.
+            prefix=f"model.{prefix}",
         )
         self.image_rope2d_emb = HunYuanRotary2DEmbedder(
             num_heads=self.num_heads,
@@ -2964,7 +2966,10 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
             seq_lens=[prefill_seq_len],
             num_image_tokens=0,
             ar_kv_reuse_len=negative_reuse_len,
-            full_attn_spans=model_kwargs["full_attn_spans"][batch_slice]
+            full_attn_spans=[
+                [(start, min(end, prefill_seq_len)) for start, end in spans if start < prefill_seq_len]
+                for spans in model_kwargs["full_attn_spans"][batch_slice]
+            ]
             if model_kwargs.get("full_attn_spans")
             else None,
         )
@@ -3019,6 +3024,28 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         device,
     ):
         ar_kv_data = model_kwargs.pop("ar_kv_data", None)
+        computed_tokens = get_paged_kv_computed_tokens()
+        if computed_tokens and computed_tokens[0] > 0:
+            positive_reuse_len = computed_tokens[0]
+            if (
+                len(computed_tokens) > 1
+                and computed_tokens[1] < positive_reuse_len
+                and (not cfg_parallel_ready or cfg_rank == 1)
+            ):
+                prefill_inputs = self._build_negative_cfg_prefill_inputs(
+                    input_ids,
+                    model_kwargs,
+                    batch_size,
+                    computed_tokens[1],
+                    positive_reuse_len,
+                    cfg_parallel_ready,
+                )
+                with (
+                    paged_kv_prefill(1, positive_reuse_len),
+                    torch.autocast(device_type=device.type, dtype=torch.bfloat16),
+                ):
+                    self.model.forward_call(**prefill_inputs)
+            return self._truncate_reused_prefix(input_ids, model_kwargs, positive_reuse_len), positive_reuse_len
         if ar_kv_data is None:
             logger.debug(
                 "[AR KV Reuse] cfg_rank=%s: no AR KV received, fallback to full recompute (reuse_len=0)",

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -367,6 +367,8 @@ class HunyuanImage3Pipeline(
 
     def __init__(self, od_config: OmniDiffusionConfig) -> None:
         self.hf_config = get_config(od_config.model, trust_remote_code=True)
+        self.hf_config.cfg_distilled = _config_flag(self.hf_config, "cfg_distilled")
+        self.hf_config.use_meanflow = _config_flag(self.hf_config, "use_meanflow")
         super().__init__(self.hf_config)
         # update diffusion config
         self.generation_config = GenerationConfig.from_pretrained(od_config.model)
@@ -1282,7 +1284,9 @@ class HunyuanImage3Pipeline(
 
         # 4. Encode conditional images
         # Skip encoding if AR KV reuse is enabled
-        has_ar_kv = kwargs.get("ar_kv_data")
+        has_ar_kv = kwargs.get("ar_kv_data") or request_layout_utils.native_kv_covers_cond_images(
+            output, kwargs.get("kv_computed_tokens", ())
+        )
         if batch_cond_image_info is not None and len(batch_cond_image_info[0]) > 0 and not has_ar_kv:
             cond_vae_images, cond_timestep, cond_vit_images = self._encode_cond_image(
                 batch_cond_image_info, cfg_factor[mode], generator=generator
@@ -1894,6 +1898,7 @@ class HunyuanImage3Pipeline(
         pipe._guidance_scale = guidance_scale
         pipe._guidance_rescale = getattr(sampling, "guidance_rescale", 0.0)
 
+        ar_kv_kwargs = self._extract_ar_kv_from_sampling(sampling)
         model_kwargs = self.prepare_model_inputs(
             prompt=prompt,
             cot_text=cot_text,
@@ -1906,8 +1911,10 @@ class HunyuanImage3Pipeline(
             batch_cond_image_info=batch_cond_image_info,
             bot_task=tokenizer_bot_task,
             prepared_layout=request_layout_utils.get_hunyuan_prepared_layout(state),
+            kv_computed_tokens=state.extra.get("kv_computed_tokens", ()),
+            **ar_kv_kwargs,
         )
-        model_kwargs.update(self._extract_ar_kv_from_sampling(sampling))
+        model_kwargs.update(ar_kv_kwargs)
         model_kwargs["use_cache"] = False
 
         input_ids = model_kwargs.pop("input_ids")
@@ -2337,6 +2344,7 @@ class HunyuanImage3Pipeline(
             batch_cond_image_info=batch_cond_image_info,
             bot_task=tokenizer_bot_task,
             prepared_layout=request_layout_utils.get_hunyuan_prepared_layout(req.requests[0]),
+            kv_computed_tokens=getattr(req.requests[0], "kv_computed_tokens", ()),
             **ar_kv_kwargs,
         )
 

vllm_omni/diffusion/models/hunyuan_image3/request_layout.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
 
@@ -373,7 +373,32 @@ def build_hunyuan_diffusion_kv_requests(
     real_pos = tokenizer_output.real_pos
     assert prefix_positions is not None and real_pos is not None
 
+    def boundary_rows(field_name: str) -> list[int | None]:
+        positions = getattr(tokenizer_output, field_name)
+        if positions is None:
+            return [None] * cfg_factor
+        if not isinstance(positions, list) or not positions:
+            raise ValueError(f"Hunyuan {field_name} must be a non-empty boundary list")
+        # Batched tokenizer output contains a separate boundary for each CFG
+        # row. Never broadcast a conditional boundary into a negative row.
+        if cfg_factor == 1 and len(positions) == 1 and (positions[0] is None or type(positions[0]) is int):
+            return positions
+        if len(positions) == cfg_factor and all(
+            isinstance(row, list) and len(row) == 1 and (row[0] is None or type(row[0]) is int) for row in positions
+        ):
+            return [row[0] for row in positions]
+        raise ValueError(f"Hunyuan {field_name} must contain one boundary per CFG row")
+
+    think_boundaries = boundary_rows("think_recaption_end_pos")
+    uncond_boundaries = boundary_rows("uncond_cfg_start_pos")
+
     target_len = hunyuan_num_image_tokens(prepared_layout.generated_image_info)
+    reusable_lens = [0] * cfg_factor
+    if think_boundaries[0] is not None:
+        reusable_lens[0] = think_boundaries[0] or 0
+    if any(boundary is not None for boundary in uncond_boundaries):
+        for sequence_id in range(1, cfg_factor):
+            reusable_lens[sequence_id] = min(reusable_lens[0], uncond_boundaries[sequence_id] or 0)
     return tuple(
         DiffusionKVRequest(
             f"{request.request_id}/diffusion-kv/{sequence_id}",
@@ -383,6 +408,7 @@ def build_hunyuan_diffusion_kv_requests(
             prefix_len=int(prefix_row[-1].item()),
             target_len=target_len,
             seq_len=int(valid_row[-1].item()),
+            prompt_token_ids=tokenizer_output.tokens[sequence_id, : reusable_lens[sequence_id]].tolist(),
             # Prompt and reference-image tokens are already embedded in this
             # row's primary self-attention sequence. Hunyuan therefore has no
             # independently projected cross/joint-attention KV context.
@@ -401,3 +427,14 @@ def get_hunyuan_prepared_layout(source: Any) -> HunyuanPreparedLayout | None:
             f"HunyuanImage3 expected prepared_layout to be HunyuanPreparedLayout, got {type(prepared_layout).__name__}"
         )
     return prepared_layout
+
+
+def native_kv_covers_cond_images(output: TokenizerEncodeOutput, computed_tokens: tuple[int, ...]) -> bool:
+    """Skip image encoding only when every CFG row already contains its image KV."""
+    image_slices = output.joint_image_slices
+    if not image_slices or len(image_slices) != len(computed_tokens):
+        return False
+    return all(
+        slices and all(image_slice.stop <= computed for image_slice in slices)
+        for slices, computed in zip(image_slices, computed_tokens, strict=True)
+    )

vllm_omni/diffusion/request.py

@@ -70,6 +70,8 @@ class OmniDiffusionRequest:
     # 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
+    # Worker-populated scheduler reuse boundaries, available before encoding.
+    kv_computed_tokens: tuple[int, ...] = ()
     # 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

@@ -11,10 +11,17 @@ from dataclasses import fields
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.outputs import KVConnectorOutput
+from vllm.v1.request import RequestStatus
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
-from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager
+from vllm_omni.diffusion.diffusion_kv.kv_connector import (
+    KVTransferRegistrationError,
+    commit_kv_load,
+    prepare_kv_requests,
+)
+from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVCacheManager
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
@@ -58,6 +65,12 @@ class BaseScheduler(ABC):
         self._prefetch_enabled: bool = False
         self._diffusion_kv_manager: DiffusionKVCacheManager | None = None
         self._kv_connector = None
+        self._kv_transfer_request_ids: set[str] = set()
+        self._kv_finished_request_ids: set[str] = set()
+        self._kv_loading_request_ids: set[str] = set()
+        self._kv_request_generations: dict[str, int] = {}
+        self._kv_draining_requests: dict[str, SchedulerRequestState] = {}
+        self._kv_received_request_ids: set[str] = set()
 
     def initialize(
         self,
@@ -75,6 +88,7 @@ class BaseScheduler(ABC):
         self._running.clear()
         self._running_sampling_params_key = None
         self._finished_req_ids.clear()
+        self._reset_kv_transfer_state()
         max_num_seqs = getattr(od_config, "max_num_seqs", 1)
         try:
             self.max_num_running_reqs = max(1, int(max_num_seqs))
@@ -114,7 +128,7 @@ class BaseScheduler(ABC):
             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._kv_connector = create_scheduler_kv_connector(od_config, kv_cache_config, kv_vllm_config)
         self._reset_scheduler_state()
 
     @property
@@ -127,7 +141,7 @@ class BaseScheduler(ABC):
         return self._add_request_with_request_id(request.request_id, request)
 
     def _add_request_with_request_id(self, request_id: str, request: OmniDiffusionRequest) -> str:
-        if request_id in self._request_states:
+        if request_id in self._request_states or request_id in self._kv_draining_requests:
             raise ValueError(f"request_id {request_id!r} is already active.")
         state = self._make_request_state(request_id, request)
         state.queued_at = time.perf_counter()
@@ -147,7 +161,10 @@ class BaseScheduler(ABC):
                 scheduled_cached_request_ids.append(request_id)
 
         # Second, schedule WAITING requests while capacity remains.
-        while self._waiting and len(self._running) < self.max_num_running_reqs:
+        deferred: list[str] = []
+        waiting_budget = len(self._waiting)
+        while self._waiting and waiting_budget and len(self._running) < self.max_num_running_reqs:
+            waiting_budget -= 1
             request_id = self._waiting[0]
             state = self._request_states.get(request_id)
             if state is None:
@@ -161,21 +178,24 @@ class BaseScheduler(ABC):
                 if self._diffusion_kv_manager.has_request(request_id):
                     diffusion_kv_metadata = self._diffusion_kv_manager.get_metadata(request_id)
                 else:
+                    matched_tokens: list[int] = []
+                    if self._kv_connector is not None:
+                        for request in state.diffusion_kv_requests:
+                            num_tokens, _ = self._kv_connector.get_num_new_matched_tokens(request, 0)
+                            if num_tokens is None:
+                                break
+                            matched_tokens.append(num_tokens)
+                        if len(matched_tokens) != len(state.diffusion_kv_requests):
+                            deferred.append(self._waiting.popleft())
+                            continue
                     try:
                         allocation = self._diffusion_kv_manager.reserve_request(
                             request_id,
                             state.diffusion_kv_requests,
                         )
                     except Exception as exc:
-                        # schedule() runs outside the Engine execution-error
-                        # wrapper. Convert both expected admission failures and
-                        # unexpected native allocator failures into a terminal
-                        # request error so the busy loop can wake its stream.
-                        if not isinstance(exc, DiffusionKVAdmissionError):
-                            logger.exception(
-                                "Unexpected Diffusion KV allocation failure for request %s",
-                                request_id,
-                            )
+                        # Reservation rolls back on failure and has not yet
+                        # handed any destination pages to the connector.
                         self._finish_requests(
                             {request_id: DiffusionRequestStatus.FINISHED_ERROR},
                             {request_id: str(exc)},
@@ -184,6 +204,27 @@ class BaseScheduler(ABC):
                     if allocation is None:
                         break
                     diffusion_kv_metadata = allocation
+                    self._kv_request_generations[request_id] = allocation.allocation_generation
+                    if self._kv_connector is not None:
+                        try:
+                            transfer_ids = commit_kv_load(
+                                self._kv_connector,
+                                self._diffusion_kv_manager.native_manager,
+                                state.diffusion_kv_requests,
+                                matched_tokens,
+                            )
+                        except KVTransferRegistrationError as exc:
+                            self._finish_requests(
+                                {request_id: DiffusionRequestStatus.FINISHED_ERROR}, {request_id: str(exc)}
+                            )
+                            continue
+                        # Other exceptions have no proven rollback. Keep pages
+                        # allocated until the Engine shuts down the Workers.
+                        self._kv_transfer_request_ids.update(transfer_ids)
+                        if transfer_ids:
+                            self._kv_loading_request_ids.add(request_id)
+                        for sequence, request in zip(allocation.sequences, state.diffusion_kv_requests, strict=True):
+                            sequence.num_computed_tokens = request.num_computed_tokens
 
             self._waiting.popleft()
             was_new_request = state.status == DiffusionRequestStatus.WAITING
@@ -202,6 +243,8 @@ class BaseScheduler(ABC):
             else:
                 scheduled_cached_request_ids.append(request_id)
 
+        self._waiting.extendleft(reversed(deferred))
+
         # Expose the next waiting request (serial mode) so the runner can
         # prefetch its KV during this forward.  Skip a request without
         # kv_sender_info (would target the wrong sender under multi-replica) or
@@ -226,11 +269,19 @@ class BaseScheduler(ABC):
             num_running_reqs=len(self._running),
             num_waiting_reqs=len(self._waiting),
             kv_prefetch_job=kv_prefetch_job,
+            kv_finished_request_ids=set(self._kv_finished_request_ids),
         )
+        if self._kv_connector is not None and (
+            self._kv_transfer_request_ids or self._kv_finished_request_ids or self._kv_draining_requests
+        ):
+            scheduler_output.kv_connector_metadata = self._kv_connector.build_connector_meta(scheduler_output)
+            scheduler_output.kv_transfer_request_ids = self._kv_transfer_request_ids
+            self._kv_transfer_request_ids = set()
 
         # update after schedule
         self._step_id += 1
         self._finished_req_ids.clear()
+        self._kv_finished_request_ids.clear()
         return scheduler_output
 
     @abstractmethod
@@ -238,7 +289,53 @@ class BaseScheduler(ABC):
         pass
 
     def has_requests(self) -> bool:
-        return bool(self._waiting or self._running)
+        return bool(
+            self._waiting
+            or self._running
+            or self._kv_transfer_request_ids
+            or self._kv_finished_request_ids
+            or self._kv_draining_requests
+        )
+
+    def update_kv_connector_output(self, output: KVConnectorOutput | None) -> None:
+        if output is None:
+            return
+        assert self._kv_connector is not None
+        self._kv_connector.update_connector_output(output)
+        finished_ids = output.finished_recving or set()
+        if not finished_ids:
+            return
+        for request_id in tuple(self._kv_loading_request_ids):
+            state = self._request_states.get(request_id) or self._kv_draining_requests.get(request_id)
+            if state is None:
+                continue
+            internal_ids = {request.request_id for request in state.diffusion_kv_requests}
+            self._kv_received_request_ids.update(internal_ids.intersection(finished_ids))
+            if internal_ids and internal_ids.issubset(self._kv_received_request_ids):
+                self._kv_loading_request_ids.discard(request_id)
+                self._kv_received_request_ids.difference_update(internal_ids)
+
+    def fail_incomplete_kv_loads(self, transfer_ids: set[str]) -> set[str]:
+        statuses = {}
+        for request_id in self._kv_loading_request_ids:
+            state = self._request_states.get(request_id)
+            if state is not None and any(req.request_id in transfer_ids for req in state.diffusion_kv_requests):
+                statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+        return self._finish_requests(statuses, {rid: "Timed out receiving diffusion KV" for rid in statuses})
+
+    def completed_kv_drains(self) -> set[str]:
+        return self._kv_draining_requests.keys() - self._kv_loading_request_ids
+
+    def release_kv_drains(self, request_ids: set[str]) -> None:
+        """Called only after all ranks completed and Worker row cleanup succeeded."""
+        for request_id in request_ids:
+            if request_id in self._kv_loading_request_ids:
+                raise RuntimeError("Cannot release diffusion pages before KV receive completes")
+            state = self._kv_draining_requests[request_id]
+            self._release_request_kv(state, state.status)
+            del self._kv_draining_requests[request_id]
+            if request_id not in self._request_states:
+                self._kv_request_generations.pop(request_id, None)
 
     def num_waiting_requests(self) -> int:
         return len(self._waiting)
@@ -270,13 +367,27 @@ class BaseScheduler(ABC):
     def get_request_state(self, request_id: str) -> SchedulerRequestState | None:
         return self._request_states.get(request_id)
 
+    def get_diffusion_kv_cleanup_targets(self, request_ids: list[str]) -> list[str | tuple[str, int]]:
+        targets: list[str | tuple[str, int]] = []
+        for request_id in request_ids:
+            if request_id in self._kv_loading_request_ids:
+                continue
+            generation = self._kv_request_generations.get(request_id)
+            if generation is not None:
+                targets.append((request_id, generation))
+            else:
+                targets.append(request_id)
+        return targets
+
     def pop_request_state(self, request_id: str) -> SchedulerRequestState | None:
-        if self._diffusion_kv_manager is not None:
-            self._diffusion_kv_manager.free_request(request_id)
         self._pop_extra_request_state(request_id)
+        if request_id not in self._kv_draining_requests:
+            self._kv_request_generations.pop(request_id, None)
         return self._request_states.pop(request_id, None)
 
     def preempt_request(self, request_id: str) -> bool:
+        if request_id in self._kv_loading_request_ids:
+            return False
         if request_id not in self._request_states:
             return False
         if request_id in self._running:
@@ -295,13 +406,7 @@ class BaseScheduler(ABC):
         self._finish_requests({request_id: status for request_id in request_ids})
 
     def close(self) -> None:
-        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._shutdown_diffusion_kv()
         self._request_states.clear()
         self._waiting.clear()
         self._running.clear()
@@ -309,6 +414,26 @@ class BaseScheduler(ABC):
         self._finished_req_ids.clear()
         self._reset_scheduler_state()
 
+    def _reset_kv_transfer_state(self) -> None:
+        """Clear request bookkeeping for native KV transfers."""
+        self._kv_transfer_request_ids.clear()
+        self._kv_finished_request_ids.clear()
+        self._kv_loading_request_ids.clear()
+        self._kv_draining_requests.clear()
+        self._kv_received_request_ids.clear()
+        self._kv_request_generations.clear()
+
+    def _shutdown_diffusion_kv(self) -> None:
+        """Stop transfers before releasing their cache and request state."""
+        from vllm_omni.diffusion.diffusion_kv.kv_connector import shutdown_kv_connector
+
+        shutdown_kv_connector(scheduler_connector=self._kv_connector)
+        self._kv_connector = None
+        if self._diffusion_kv_manager is not None:
+            self._diffusion_kv_manager.close()
+            self._diffusion_kv_manager = None
+        self._reset_kv_transfer_state()
+
     def _finish_requests(
         self,
         statuses: dict[str, DiffusionRequestStatus],
@@ -326,7 +451,6 @@ class BaseScheduler(ABC):
             state = self._request_states.get(request_id)
             if state is None or state.is_finished():
                 continue
-
             finished_req_ids.add(request_id)
             if request_id in self._running:
                 running_to_remove.add(request_id)
@@ -341,10 +465,12 @@ class BaseScheduler(ABC):
             self._waiting = deque(request_id for request_id in self._waiting if request_id not in waiting_to_remove)
 
         for request_id in finished_req_ids:
-            if self._diffusion_kv_manager is not None:
-                self._diffusion_kv_manager.free_request(request_id)
             state = self._request_states[request_id]
             status = statuses[request_id]
+            if request_id in self._kv_loading_request_ids:
+                self._kv_draining_requests[request_id] = state
+            else:
+                self._release_request_kv(state, status)
             state.status = status
             if status == DiffusionRequestStatus.FINISHED_ERROR:
                 state.error = None if errors is None else errors.get(request_id)
@@ -354,6 +480,27 @@ class BaseScheduler(ABC):
         self._finished_req_ids |= finished_req_ids
         return finished_req_ids
 
+    def _release_request_kv(self, state: SchedulerRequestState, status: DiffusionRequestStatus) -> None:
+        request_id = state.request_id
+        if self._kv_connector is not None:
+            for request in state.diffusion_kv_requests:
+                request.status = (
+                    RequestStatus.FINISHED_ABORTED
+                    if status == DiffusionRequestStatus.FINISHED_ABORTED
+                    else RequestStatus.FINISHED_STOPPED
+                )
+                if request.kv_transfer_params:
+                    self._kv_finished_request_ids.add(request.request_id)
+                block_ids = []
+                if self._diffusion_kv_manager.has_request(request_id):
+                    blocks = self._diffusion_kv_manager.native_manager.get_blocks(request.request_id)
+                    block_ids = blocks.get_block_ids()[0]
+                delay_free, _ = self._kv_connector.request_finished(request, block_ids)
+                if delay_free:
+                    raise RuntimeError("Diffusion consumer cannot release pages while KV transfer is active")
+        if self._diffusion_kv_manager is not None:
+            self._diffusion_kv_manager.free_request(request_id)
+
     def _finalize_update_from_output(
         self,
         sched_output: DiffusionSchedulerOutput,
@@ -391,13 +538,9 @@ 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)
+            prepare_kv_requests(kv_requests, kv_transfer_params)
 
         # DiffusionKVRequest objects are mutable Scheduler/native-KVCacheManager
         # state and must never ride the normal request payload to a Worker.

vllm_omni/diffusion/sched/interface.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import enum
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, TypedDict
 
@@ -254,9 +254,10 @@ 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
+    kv_transfer_request_ids: set[str] = field(default_factory=set)
+    # Connector lifecycle uses per-sequence IDs, not public request IDs.
+    kv_finished_request_ids: set[str] = field(default_factory=set)
 
     @cached_property
     def scheduled_request_ids(self) -> list[str]:
@@ -274,4 +275,4 @@ class DiffusionSchedulerOutput:
 
     @property
     def is_empty(self) -> bool:
-        return self.num_scheduled_reqs == 0
+        return self.num_scheduled_reqs == 0 and self.kv_connector_metadata is None

vllm_omni/diffusion/sched/request_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
 
@@ -100,7 +100,7 @@ class RequestScheduler(BaseScheduler):
 
     def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
         scheduled_request_ids = sched_output.scheduled_request_ids
-        if not scheduled_request_ids:
+        if not scheduled_request_ids and not sched_output.finished_req_ids:
             return set()
 
         terminal_statuses: dict[str, DiffusionRequestStatus] = {}

vllm_omni/diffusion/sched/step_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
 
@@ -64,7 +64,7 @@ class StepScheduler(BaseScheduler):
 
     def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
         scheduled_request_ids = sched_output.scheduled_request_ids
-        if not scheduled_request_ids:
+        if not scheduled_request_ids and not sched_output.finished_req_ids:
             return set()
 
         terminal_statuses: dict[str, DiffusionRequestStatus] = {}

vllm_omni/diffusion/vllm_config.py

@@ -151,6 +151,9 @@ class _DiffusionVllmModelConfig:
     def get_head_size(self) -> int:
         return self._require_attention_geometry()[2]
 
+    def get_total_num_kv_heads(self) -> int:
+        return self.hf_text_config.num_key_value_heads
+
     @property
     def is_diffusion(self) -> bool:
         # This config is also installed as the current VllmConfig for the
@@ -224,7 +227,10 @@ def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiff
         vllm_config.cache_config.kv_cache_memory_bytes = getattr(od_config, "kv_cache_memory_bytes", None)
         # Prefix hashes/publication are intentionally deferred to the next PR.
         vllm_config.cache_config.enable_prefix_caching = False
-        vllm_config.scheduler_config.max_num_seqs = int(od_config.max_num_seqs)
+        max_rows_per_request = getattr(od_config, "diffusion_kv_max_rows_per_request", None)
+        if type(max_rows_per_request) is not int or max_rows_per_request <= 0:
+            raise ValueError("paged_scheduler requires a positive diffusion_kv_max_rows_per_request")
+        vllm_config.scheduler_config.max_num_seqs = int(od_config.max_num_seqs) * max_rows_per_request
         max_num_batched_tokens = getattr(od_config, "max_num_batched_tokens", None)
         if max_num_batched_tokens is not None:
             vllm_config.scheduler_config.max_num_batched_tokens = int(max_num_batched_tokens)

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -193,14 +193,20 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         self.state_cache: dict[str, StepRequestState] = {}
 
         # Initialize KV cache manager for connector management.
-        self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config)
+        self.kv_transfer_manager = (
+            OmniKVTransferManager.from_od_config(od_config)
+            if getattr(od_config, "kv_transfer_config", None) is None
+            else None
+        )
+        self._kv_connector = None
 
         # Prefetch covers TP / SP / CFG-Parallel / HSDP.  Disabled when a CFG
         # companion KV collector is set (that KV is not backgrounded).
         has_cfg_companion_kv = getattr(od_config, "cfg_kv_collect_func", None) is not None
 
         self._kv_prefetch_enabled = (
-            bool(self.kv_transfer_manager.config.enable_kv_async_prefetch)
+            self.kv_transfer_manager is not None
+            and bool(self.kv_transfer_manager.config.enable_kv_async_prefetch)
             and not has_cfg_companion_kv
             and self.kv_transfer_manager.config.need_recv_cache
         )
@@ -471,6 +477,16 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
         self.diffusion_kv_backend.initialize_kv_cache(kv_cache_config)
         self.kv_cache_config = self.diffusion_kv_backend.kv_cache_config
+        if getattr(self.od_config, "kv_transfer_config", None) is not None:
+            from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
+
+            self._kv_connector = ActiveKVConnector(self.vllm_config, self.diffusion_kv_backend.kv_caches_by_layer)
+
+    def prepare_kv_for_forward(self, scheduler_output: DiffusionSchedulerOutput):
+        from vllm_omni.diffusion.diffusion_kv.kv_connector import wait_for_kv_load
+
+        timeout = self.od_config.kv_transfer_config.kv_connector_extra_config.get("transfer_timeout", 60.0)
+        return wait_for_kv_load(self._kv_connector, scheduler_output, timeout)
 
     def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool:
         return self.diffusion_kv_backend.install_diffusion_kv_metadata(metadata)
@@ -483,7 +499,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
     ) -> int:
         return self.diffusion_kv_backend.get_diffusion_kv_row(request_id, sequence_id, context_id)
 
-    def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int:
+    def remove_diffusion_kv_requests(self, request_ids: list[str | tuple[str, int]]) -> int:
         return self.diffusion_kv_backend.remove_diffusion_kv_requests(request_ids)
 
     def refresh_diffusion_kv_block_table_layout(self) -> None:
@@ -520,8 +536,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     DiffusionPagedAttentionRow(
                         request_id=request_metadata.request_id,
                         sequence_id=sequence.sequence_id,
-                        query_len=sequence.seq_len,
+                        query_len=sequence.seq_len - sequence.num_computed_tokens,
                         seq_len=sequence.seq_len,
+                        kv_start_pos=sequence.num_computed_tokens,
                     )
                 )
                 denoise_rows.append(
@@ -606,6 +623,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         kv_prefetch_job: KVPrefetchJob | None = None,
         use_prefetch: bool = False,
     ) -> None:
+        if self.kv_transfer_manager is None:
+            self._initialize_generator(req.sampling_params)
+            return
         # Receive AR KV. Single-request execution can use the prefetch path:
         # consume prior-forward payload, sync-fallback on miss; request-batch
         # execution keeps the synchronous per-request receive path.
@@ -770,6 +790,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                         "Diffusion KV metadata count must match the request batch: "
                         f"metadata={len(diffusion_kv_metadata)}, requests={len(reqs)}"
                     )
+                for req, metadata in zip(reqs, diffusion_kv_metadata, strict=True):
+                    req.kv_computed_tokens = tuple(seq.num_computed_tokens for seq in metadata.sequences)
                 paged_metadata = self._build_paged_attention_metadata(diffusion_kv_metadata)
                 paged_kv_runtime, paged_kv_context = self.diffusion_kv_backend.activate_paged_attention_metadata(
                     paged_metadata
@@ -1018,13 +1040,18 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     kv_sender_info=sched_new_req.req.kv_sender_info,
                     prepared_layout=getattr(sched_new_req.req, "prepared_layout", None),
                 )
+                if sched_new_req.diffusion_kv_metadata is not None:
+                    new_state.extra["kv_computed_tokens"] = tuple(
+                        seq.num_computed_tokens for seq in sched_new_req.diffusion_kv_metadata.sequences
+                    )
                 state_req = copy.copy(sched_new_req.req)
                 state_req.sampling_params = new_state.sampling
-                self.kv_transfer_manager.receive_multi_kv_cache_distributed(
-                    state_req,
-                    cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
-                    target_device=self._target_device,
-                )
+                if self.kv_transfer_manager is not None:
+                    self.kv_transfer_manager.receive_multi_kv_cache_distributed(
+                        state_req,
+                        cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
+                        target_device=self._target_device,
+                    )
                 self.state_cache[request_id] = new_state
                 resolved.append(new_state)
 

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -59,6 +59,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     get_hsdp_replicate_group,
     get_pp_group,
     get_sp_group,
+    get_world_group,
     init_distributed_environment,
     initialize_model_parallel,
     model_parallel_is_initialized,
@@ -116,7 +117,11 @@ def _all_gather_rank_values(value: Any) -> list[Any]:
     if not dist.is_available() or not dist.is_initialized():
         return [value]
     values: list[Any] = [None] * dist.get_world_size()
-    dist.all_gather_object(values, value)
+    # Object collectives are control-plane traffic. Using the default NCCL
+    # group serializes them through temporary CUDA tensors, adding pointless
+    # H2D/DtoH copies to every rank-wide status check. Reuse the world
+    # coordinator's Gloo group, as vLLM does for CPU metadata collectives.
+    dist.all_gather_object(values, value, group=get_world_group().cpu_group)
     return values
 
 
@@ -446,15 +451,13 @@ class DiffusionWorker:
                 raise RuntimeError("Diffusion KV memory snapshot was not captured before model loading")
             override = self.vllm_config.cache_config.kv_cache_memory_bytes
             if override:
-                # Match native vLLM: an explicit cache budget skips automatic
-                # capacity derivation, but still runs the maximum-shape model
-                # request so lazy kernels and communication buffers initialize.
-                self.model_runner.profile_run(profile_requests)
+                # Post-allocation warmup avoids retaining a second activation arena.
                 logger.info(
                     "Worker %d: Initial free memory %s GiB, reserved %s GiB memory for "
                     "Diffusion KV Cache as specified by kv_cache_memory_bytes config and "
                     "skipped automatic memory profiling. This does not respect the "
-                    "gpu_memory_utilization config. A profile warmup was still executed.",
+                    "gpu_memory_utilization config. Model kernels will be initialized by "
+                    "the post-allocation startup warmup.",
                     self.rank,
                     format_gib(self.init_snapshot.free_memory),
                     format_gib(int(override)),
@@ -467,6 +470,8 @@ class DiffusionWorker:
             ) as profile_result:
                 self.model_runner.profile_run(profile_requests)
 
+            current_omni_platform.empty_cache()
+
             available_memory = self.requested_memory - profile_result.non_kv_cache_memory
             if available_memory <= 0:
                 raise RuntimeError(
@@ -534,12 +539,18 @@ class DiffusionWorker:
         with self._maybe_get_memory_pool_context("kv_cache"):
             self.model_runner.set_kv_cache_config(kv_cache_config)
 
-    def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int:
+    def remove_diffusion_kv_requests(self, request_ids: list[str | tuple[str, int]]) -> int:
         """Clear Worker-local rows without freeing Scheduler-owned blocks."""
 
         assert self.model_runner is not None, "Model runner not initialized"
         return self.model_runner.remove_diffusion_kv_requests(request_ids)
 
+    def prepare_kv_for_forward(self, scheduler_output: DiffusionSchedulerOutput):
+        return _run_and_gather_rank_values(
+            "Diffusion KV receive",
+            lambda: self.model_runner.prepare_kv_for_forward(scheduler_output),
+        )
+
     def init_lora_manager(self) -> None:
         """Initialize the LoRA manager for this worker."""
         pipeline = self.model_runner.pipeline
@@ -790,6 +801,15 @@ class DiffusionWorker:
         Args:
             level: Sleep level. Level 1 offloads weights, level 2 also saves buffers.
         """
+        # The config validator rejects sleep for the native paged path. Keep
+        # this worker-side guard precise as well: test doubles and legacy
+        # configs may expose arbitrary attributes through Mock/getattr.
+        if (
+            getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            is DiffusionKVCacheMode.PAGED_SCHEDULER
+            and getattr(self.od_config, "kv_transfer_config", None) is not None
+        ):
+            raise ValueError("Cannot sleep while native KV connector memory is registered")
         CuMemAllocator = _get_cumem_allocator_class()
         allocator = CuMemAllocator.get_instance()
 
@@ -1254,7 +1274,11 @@ class WorkerProc:
 
         world_size = torch.distributed.get_world_size()
         statuses: list[dict[str, Any] | None] = [None] * world_size
-        torch.distributed.all_gather_object(statuses, status)
+        torch.distributed.all_gather_object(
+            statuses,
+            status,
+            group=get_world_group().cpu_group,
+        )
         missing_ranks = [rank for rank, rank_status in enumerate(statuses) if rank_status is None]
         if missing_ranks:
             logger.warning("RPC rank status gather returned missing entries for ranks: %s", missing_ranks)

vllm_omni/engine/errors.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Engine runtime error types."""
+
+
+class NativeKVHandoffError(RuntimeError):
+    """One request lost its producer binding or transfer metadata."""

vllm_omni/engine/orchestrator.py

@@ -38,6 +38,7 @@ from vllm_omni.diffusion.data import is_diffusion_request_started_output
 from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
+from vllm_omni.engine.errors import NativeKVHandoffError
 from vllm_omni.engine.membership_controller import MembershipController
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
@@ -211,6 +212,7 @@ class OrchestratorRequestState:
     session_owned: bool = False
     running_counter_registered: bool = False
     request_artifact_dirs: set[str] = field(default_factory=set)
+    native_kv_transfer_id: str | None = None
 
 
 @dataclass
@@ -1363,6 +1365,16 @@ class OrchestratorBase:
         try:
             await dispatch()
             return True
+        except NativeKVHandoffError as e:
+            await self._fail_request_client_error(
+                req_id,
+                stage_id,
+                str(e),
+                status_code=HTTPStatus.BAD_GATEWAY.value,
+                error_type="NativeKVHandoffError",
+                release_owners=True,
+            )
+            return False
         except StageUnavailableError as e:
             # No specific replica to evict: the stage already has no live
             # replica or the chosen slot was evicted. Fail just this request.
@@ -2193,12 +2205,13 @@ class OrchestratorBase:
                     req_id,
                     req_state,
                     diffusion_prompt,
-                    submit_kwargs={
-                        "kv_sender_info": self._build_kv_sender_info(
-                            list(getattr(next_client, "engine_input_source", None) or [src_stage_id]),
-                            request_id=req_id,
-                        )
-                    },
+                    submit_kwargs=self._diffusion_submit_kwargs(
+                        req_id,
+                        src_stage_id,
+                        next_client,
+                        req_state,
+                        output,
+                    ),
                     params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params),
                 )
             self._on_stage_submitted(
@@ -2452,12 +2465,12 @@ class OrchestratorBase:
             _t_submit_start = _time.perf_counter()
 
             if next_pool.stage_type == "diffusion":
-                submit_kwargs = {
-                    "kv_sender_info": self._build_kv_sender_info(
-                        list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]),
-                        request_id=request_id,
-                    )
-                }
+                submit_kwargs = self._diffusion_submit_kwargs(
+                    request_id,
+                    next_stage_id - 1,
+                    next_pool.stage_client,
+                    req_state,
+                )
                 submitted = await self._dispatch_or_fail_request(
                     lambda: next_pool.submit_initial(
                         request_id,
@@ -2554,6 +2567,67 @@ class OrchestratorBase:
 
         return True
 
+    def _maybe_attach_native_kv_transfer_params(
+        self,
+        req_state: OrchestratorRequestState,
+        prompt: Any,
+    ) -> None:
+        source_stage_id = 0
+        if source_stage_id + 1 > req_state.final_stage_id:
+            return
+        config = getattr(self.stage_pools[source_stage_id].stage_vllm_config, "kv_transfer_config", None)
+        if getattr(config, "kv_role", None) != "kv_producer":
+            return
+        if self.stage_pools[source_stage_id + 1].stage_type != "diffusion":
+            return
+
+        from vllm_omni.diffusion.diffusion_kv.kv_connector import build_source_kv_transfer_params, mint_transfer_id
+
+        transfer_id = mint_transfer_id(req_state.request_id)
+        req_state.native_kv_transfer_id = transfer_id
+        params = build_source_kv_transfer_params(
+            transfer_id=transfer_id,
+            remote_engine_id=None,
+            remote_bootstrap_addr=None,
+        )
+        for sampling in (req_state.sampling_params_list[source_stage_id], prompt.sampling_params):
+            sampling.extra_args = {**(sampling.extra_args or {}), "kv_transfer_params": params}
+
+    def _diffusion_submit_kwargs(
+        self,
+        request_id: str,
+        source_stage_id: int,
+        next_client: Any,
+        req_state: OrchestratorRequestState,
+        output: Any = None,
+    ) -> dict[str, Any]:
+        source_stage_ids = list(getattr(next_client, "engine_input_source", None) or [source_stage_id])
+        if req_state.native_kv_transfer_id is None:
+            return {"kv_sender_info": self._build_kv_sender_info(source_stage_ids, request_id=request_id)}
+        if output is None:
+            raise NativeKVHandoffError("Native KV handoff requires one completed AR source")
+        from vllm_omni.diffusion.diffusion_kv.kv_connector import (
+            bootstrap_addr_from_kv_transfer_config,
+            build_target_kv_transfer_params,
+        )
+
+        source = self.stage_pools[source_stage_id].get_bound_client(request_id)
+        config = getattr(getattr(source, "vllm_config", None), "kv_transfer_config", None)
+        if source is None or config is None:
+            raise NativeKVHandoffError(
+                f"Native KV handoff for {request_id}: bound AR replica or its KV configuration is unavailable"
+            )
+        params = getattr(output, "kv_transfer_params", None)
+        if not params or "num_transfer_tokens" not in params:
+            raise NativeKVHandoffError("AR source completed without native KV transfer metadata")
+        return {
+            "kv_transfer_params": build_target_kv_transfer_params(
+                source_params=params,
+                remote_engine_id=config.engine_id,
+                remote_bootstrap_addr=bootstrap_addr_from_kv_transfer_config(config),
+            )
+        }
+
     def _build_kv_sender_info(
         self,
         sender_stage_ids: list[int],
@@ -2677,6 +2751,7 @@ class Orchestrator(OrchestratorBase):
             request_artifact_dirs=set(msg.request_artifact_dirs or ()),
         )
         self.request_states[request_id] = req_state
+        self._maybe_attach_native_kv_transfer_params(req_state, prompt)
         self._register_running_request(req_state)
         req_state.streaming.enabled = bool(getattr(prompt, "resumable", False))
         req_state.stage_submit_ts[stage_id] = _time.time()

vllm_omni/engine/stage_init_utils.py

@@ -1189,6 +1189,8 @@ def _project_omni_stage_engine_args(
     engine_args["async_chunk"] = connector_config.async_chunk
     if connector_config.omni_kv_config is not None:
         engine_args["omni_kv_config"] = copy.deepcopy(connector_config.omni_kv_config)
+    if connector_config.kv_transfer_config is not None:
+        engine_args["kv_transfer_config"] = copy.deepcopy(connector_config.kv_transfer_config)
 
     if is_diffusion:
         engine_args["parallel_config"] = _project_omni_config_fields(
@@ -1980,6 +1982,11 @@ def build_diffusion_config(
     )
     od_config = OmniDiffusionConfig(**{name: value for name, value in diffusion_kwargs.items() if value is not None})
 
+    for dimension in ("height", "width"):
+        value = getattr(metadata.default_sampling_params, dimension, None)
+        if isinstance(value, int) and value > 0:
+            od_config.additional_config.setdefault(f"diffusion_kv_profile_{dimension}", value)
+
     num_devices_per_stage = od_config.parallel_config.world_size
     device_control_env = current_omni_platform.device_control_env_var
     visible_devices_str = os.environ.get(device_control_env) if device_control_env else None

vllm_omni/engine/stage_runtime.py

@@ -19,7 +19,7 @@ from typing import Any, TypedDict, cast
 import janus
 from vllm.logger import init_logger
 
-from vllm_omni.config.omni_config import BaseVllmOmniStageConfig
+from vllm_omni.config.omni_config import BaseVllmOmniStageConfig, VllmOmniDiffusionStageConfig
 from vllm_omni.distributed.omni_connectors.utils.initialization import (
     resolve_omni_kv_config_for_stage,
 )
@@ -47,6 +47,7 @@ from vllm_omni.engine.stage_engine_startup import (
 from vllm_omni.engine.stage_init_utils import (
     LogicalStageInitPlan,
     ReplicaInitPlan,
+    StageMetadata,
     _inject_inferred_kv_tp_topology,
     acquire_device_locks,
     build_engine_args_dict,
@@ -537,6 +538,7 @@ class StageRuntime:
             replicas_per_stage,
             replica_devices_map,
         )
+        self._validate_native_kv_topology(stage_plans)
         return stage_plans
 
     def _finalize_initialized_stages(
@@ -585,6 +587,113 @@ class StageRuntime:
 
     # ---- Internal methods ----
 
+    def _validate_native_kv_topology(self, plans: list[LogicalStageInitPlan]) -> None:
+        """Reject unsupported native AR->DiT graphs before launching any Worker."""
+        if not any(plan.replicas[0].metadata.stage_type == "diffusion" for plan in plans):
+            return  # Native LLM prefill/decode is outside this path.
+        roles = {}
+        for plan in plans:
+            replica = plan.replicas[0]
+            if replica.stage_vllm_config is not None:
+                config = getattr(replica.stage_vllm_config, "kv_transfer_config", None)
+            elif isinstance(replica.stage_cfg, VllmOmniDiffusionStageConfig):
+                config = (
+                    replica.stage_cfg.connector_config.kv_transfer_config
+                    or replica.stage_cfg.diffusion_config.kv_transfer_config
+                )
+            else:
+                args = getattr(replica.stage_cfg, "engine_args", {})
+                config = (
+                    args.get("kv_transfer_config")
+                    if isinstance(args, Mapping)
+                    else getattr(args, "kv_transfer_config", None)
+                )
+            if config is not None:
+                roles[plan.stage_id] = config.get("kv_role") if isinstance(config, Mapping) else config.kv_role
+        if not roles:
+            return
+        if (
+            self._async_chunk
+            or len(plans) != 2
+            or roles != {0: "kv_producer", 1: "kv_consumer"}
+            or plans[0].replicas[0].metadata.stage_type != "llm"
+            or plans[1].replicas[0].metadata.stage_type != "diffusion"
+            or list(plans[1].replicas[0].metadata.engine_input_source or []) != [0]
+        ):
+            raise ValueError(
+                "Native AR-to-DiT KV transfer currently supports only a two-stage "
+                "0 (kv_producer) -> 1 (diffusion kv_consumer, engine_input_source=[0]) "
+                "pipeline with async_chunk=False"
+            )
+
+    @staticmethod
+    def _prepare_replica_stage_config(
+        stage_cfg: Any,
+        *,
+        stage_id: int,
+        stage_type: str,
+        replica_id: int,
+        num_replicas: int,
+    ) -> tuple[Any, bool]:
+        """Isolate replica config and assign the native DiT KV identity."""
+        if isinstance(stage_cfg, BaseVllmOmniStageConfig):
+            native_kv = stage_cfg.connector_config.kv_transfer_config
+            if native_kv is None and isinstance(stage_cfg, VllmOmniDiffusionStageConfig):
+                native_kv = stage_cfg.diffusion_config.kv_transfer_config
+        else:
+            native_kv = (
+                stage_cfg.engine_args.get("kv_transfer_config")
+                if isinstance(stage_cfg.engine_args, Mapping)
+                else getattr(stage_cfg.engine_args, "kv_transfer_config", None)
+            )
+        # Keep the logical stage's device pool and native KV identity
+        # intact; each replica owns its config throughout startup.
+        replica_cfg = copy.deepcopy(stage_cfg) if num_replicas > 1 or native_kv else stage_cfg
+        if native_kv and stage_type == "diffusion":
+            if isinstance(replica_cfg, VllmOmniDiffusionStageConfig):
+                kv_config = (
+                    replica_cfg.connector_config.kv_transfer_config or replica_cfg.diffusion_config.kv_transfer_config
+                )
+                assert kv_config is not None
+                kv_config.engine_id = f"{kv_config.engine_id}-s{stage_id}-r{replica_id}"
+            else:
+                kv_config = replica_cfg.engine_args["kv_transfer_config"]
+                kv_config["engine_id"] = f"{kv_config['engine_id']}-s{stage_id}-r{replica_id}"
+        return replica_cfg, bool(native_kv)
+
+    @staticmethod
+    def _prepare_replica_vllm_config(
+        stage_vllm_config: Any,
+        replica_cfg: Any,
+        replica_metadata: StageMetadata,
+        *,
+        native_kv: bool,
+    ) -> Any:
+        """Assign native AR KV identity and the producer bootstrap endpoint."""
+        if not native_kv or stage_vllm_config is None:
+            return stage_vllm_config
+        replica_vllm_config = copy.deepcopy(stage_vllm_config)
+        kv_config = replica_vllm_config.kv_transfer_config
+        kv_config.engine_id = f"{kv_config.engine_id}-s{replica_metadata.stage_id}-r{replica_metadata.replica_id}"
+        if kv_config.kv_connector == "MooncakeConnector" and kv_config.kv_role == "kv_producer":
+            extra = kv_config.kv_connector_extra_config
+            port = int(extra.get("bootstrap_port", 8998)) + replica_metadata.replica_id
+            extra["bootstrap_addr"] = f"http://{kv_config.kv_ip}:{port}"
+            if isinstance(replica_cfg, BaseVllmOmniStageConfig):
+                runtime_cfg = replica_cfg.runtime_config
+                runtime_cfg.env = {
+                    **(runtime_cfg.env or {}),
+                    "VLLM_MOONCAKE_BOOTSTRAP_PORT": str(port),
+                }
+            else:
+                runtime_cfg = copy.deepcopy(replica_metadata.runtime_cfg or {})
+                runtime_cfg["env"] = {
+                    **(runtime_cfg.get("env") or {}),
+                    "VLLM_MOONCAKE_BOOTSTRAP_PORT": str(port),
+                }
+                replica_metadata.runtime_cfg = runtime_cfg
+        return replica_vllm_config
+
     def _build_logical_stage_init_plans(
         self,
         omni_transfer_config: Any,
@@ -658,9 +767,13 @@ class StageRuntime:
                 )
 
             for replica_id in range(num_replicas):
-                # Keep the logical stage's device pool intact; each replica owns
-                # the same config throughout planning, metadata and launch.
-                replica_cfg = copy.deepcopy(stage_cfg) if num_replicas > 1 else stage_cfg
+                replica_cfg, native_kv = self._prepare_replica_stage_config(
+                    stage_cfg,
+                    stage_id=stage_id,
+                    stage_type=base_metadata.stage_type,
+                    replica_id=replica_id,
+                    num_replicas=num_replicas,
+                )
                 if stage_idx in replica_devices_map:
                     devices = replica_devices_map[stage_idx][replica_id]
                     runtime_cfg = getattr(replica_cfg, "runtime_config", getattr(replica_cfg, "runtime", None))
@@ -673,6 +786,12 @@ class StageRuntime:
                     else extract_legacy_stage_metadata(replica_cfg)
                 )
                 replica_metadata.replica_id = replica_id
+                replica_vllm_config = self._prepare_replica_vllm_config(
+                    stage_vllm_config,
+                    replica_cfg,
+                    replica_metadata,
+                    native_kv=native_kv,
+                )
                 if launch_mode == "remote" and replica_metadata.stage_type != "diffusion":
                     replica_metadata.runtime_cfg = None
                 replicas.append(
@@ -684,7 +803,7 @@ class StageRuntime:
                         metadata=replica_metadata,
                         stage_connector_spec=stage_connector_spec,
                         omni_kv_connector=omni_kv_connector,
-                        stage_vllm_config=stage_vllm_config,
+                        stage_vllm_config=replica_vllm_config,
                         executor_class=executor_class,
                         engine_args_dict=copy.deepcopy(engine_args_dict) if engine_args_dict is not None else None,
                     )

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -1536,7 +1536,10 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         # Used to embed timestep information into the input sequence.
         self.timestep_emb = TimestepEmbedder(hidden_size=config.hidden_size)
 
-        tokenizer = get_tokenizer(vllm_config.model_config.tokenizer)
+        tokenizer = get_tokenizer(
+            vllm_config.model_config.tokenizer,
+            trust_remote_code=vllm_config.model_config.trust_remote_code,
+        )
         self._mrope_img_token_id = tokenizer.convert_tokens_to_ids("<img>")
         self._mrope_boi_token_id = tokenizer.convert_tokens_to_ids("<boi>")
         self._mrope_eoi_token_id = tokenizer.convert_tokens_to_ids("<eoi>")

vllm_omni/worker/mixins.py

@@ -7,7 +7,7 @@ from typing import Any
 
 
 class OmniWorkerMixin:
-    """Mixin to ensure Omni plugins are loaded in worker processes."""
+    """Shared Omni plugin and native KV connector setup for workers."""
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
@@ -15,6 +15,16 @@ class OmniWorkerMixin:
 
         load_omni_general_plugins()
 
+    def initialize_from_config(self, kv_cache_config) -> None:
+        # The concrete GPU/NPU/XPU worker supplies this method through the MRO.
+        super().initialize_from_config(kv_cache_config)  # type: ignore[misc]
+        from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
+
+        if has_kv_transfer_group():
+            from vllm_omni.diffusion.diffusion_kv.kv_connector import install_mooncake_cfg_fanout
+
+            install_mooncake_cfg_fanout(get_kv_transfer_group())
+
 
 def maybe_unpad_input_ids(model, input_ids, num_tokens_unpadded):
     """Trim cudagraph-bucket padding for models that split a flat ``input_ids``