[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>

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"}