diff --git a/docs/.nav.yml b/docs/.nav.yml
index beb15bf47..da55af813 100644
--- a/docs/.nav.yml
+++ b/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
diff --git a/docs/design/feature/mooncake-kv-transfer.md b/docs/design/feature/mooncake-kv-transfer.md
new file mode 100644
index 000000000..ed83d562f
--- /dev/null
+++ b/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.
diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index 76bf3c319..cdb303316 100644
--- a/tests/config/test_omni_config.py
+++ b/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",
diff --git a/tests/core/sched/test_omni_ar_scheduler_free_request_cleanup.py b/tests/core/sched/test_omni_ar_scheduler_free_request_cleanup.py
index 73c71fd63..8bc63795f 100644
--- a/tests/core/sched/test_omni_ar_scheduler_free_request_cleanup.py
+++ b/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]
diff --git a/tests/diffusion/diffusion_kv/test_block_tables.py b/tests/diffusion/diffusion_kv/test_block_tables.py
index aa8e07307..dbca89735 100644
--- a/tests/diffusion/diffusion_kv/test_block_tables.py
+++ b/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)
diff --git a/tests/diffusion/diffusion_kv/test_initialization.py b/tests/diffusion/diffusion_kv/test_initialization.py
index b5f914b11..b9161ba92 100644
--- a/tests/diffusion/diffusion_kv/test_initialization.py
+++ b/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,
)
diff --git a/tests/diffusion/diffusion_kv/test_kv_connector.py b/tests/diffusion/diffusion_kv/test_kv_connector.py
index 19847b3d5..7528a201e 100644
--- a/tests/diffusion/diffusion_kv/test_kv_connector.py
+++ b/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
diff --git a/tests/diffusion/diffusion_kv/test_metadata.py b/tests/diffusion/diffusion_kv/test_metadata.py
index 07e3eeb70..da9e7b331 100644
--- a/tests/diffusion/diffusion_kv/test_metadata.py
+++ b/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:
diff --git a/tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py b/tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py
new file mode 100644
index 000000000..d9e66ddea
--- /dev/null
+++ b/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
diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py
index 535c1b411..271f55ee8 100644
--- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py
+++ b/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
diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py
index d890dcc7b..285fb2c96 100644
--- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py
+++ b/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
diff --git a/tests/diffusion/diffusion_kv/test_spec_discovery.py b/tests/diffusion/diffusion_kv/test_spec_discovery.py
index 1ccef7a9f..f3db266a3 100644
--- a/tests/diffusion/diffusion_kv/test_spec_discovery.py
+++ b/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"}
diff --git a/tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py b/tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py
index 2a17cd691..af9a41dfc 100644
--- a/tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py
+++ b/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
diff --git a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py
index 65fc0e67b..db48175c3 100644
--- a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py
+++ b/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
diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py
index 119939ae2..bb64e5f01 100644
--- a/tests/diffusion/test_diffusion_engine_rpc_routing.py
+++ b/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,
diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py
index 69a9e4395..f5e4aaa2c 100644
--- a/tests/diffusion/test_diffusion_model_runner.py
+++ b/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=[]),
)
diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py
index 20a2f0384..38da458b0 100644
--- a/tests/diffusion/test_diffusion_scheduler.py
+++ b/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)
diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py
index ade8d8fa7..50fca3ea8 100644
--- a/tests/diffusion/test_multiproc_engine_concurrency.py
+++ b/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,
diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py
index 75301ee0a..ec897bd7d 100644
--- a/tests/engine/test_async_omni_engine_stage_init.py
+++ b/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)
diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py
index 3b6ee7e7b..775eeb86b 100644
--- a/tests/engine/test_orchestrator.py
+++ b/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:
diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py
index ed886d026..1d106932d 100644
--- a/tests/entrypoints/test_async_omni_diffusion_config.py
+++ b/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)
diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py
index 8f7319b13..ef48603cf 100644
--- a/vllm_omni/config/omni_config.py
+++ b/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,
)
diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py
index 3d866fda7..385e180af 100644
--- a/vllm_omni/core/sched/omni_ar_scheduler.py
+++ b/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
diff --git a/vllm_omni/core/sched/utils.py b/vllm_omni/core/sched/utils.py
index 7e0c16bd4..12cadae73 100644
--- a/vllm_omni/core/sched/utils.py
+++ b/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))
diff --git a/vllm_omni/deploy/hunyuan-image-3-mooncakeconnector.yaml b/vllm_omni/deploy/hunyuan-image-3-mooncakeconnector.yaml
new file mode 100644
index 000000000..ea088f916
--- /dev/null
+++ b/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
diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py
index 073fd3bd5..679bf4955 100644
--- a/vllm_omni/diffusion/attention/backends/flash_attn.py
+++ b/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(
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index e19470fa5..9dd69502f 100644
--- a/vllm_omni/diffusion/data.py
+++ b/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)
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 6e53b5909..5e3082f45 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/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,
diff --git a/vllm_omni/diffusion/diffusion_kv/kv_connector.py b/vllm_omni/diffusion/diffusion_kv/kv_connector.py
index 47c609860..eaa478300 100644
--- a/vllm_omni/diffusion/diffusion_kv/kv_connector.py
+++ b/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"):
diff --git a/vllm_omni/diffusion/diffusion_kv/manager.py b/vllm_omni/diffusion/diffusion_kv/manager.py
index eee51b864..3b2fb3b42 100644
--- a/vllm_omni/diffusion/diffusion_kv/manager.py
+++ b/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)
diff --git a/vllm_omni/diffusion/diffusion_kv/metadata.py b/vllm_omni/diffusion/diffusion_kv/metadata.py
index db6cb54ec..19ecd1dc3 100644
--- a/vllm_omni/diffusion/diffusion_kv/metadata.py
+++ b/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
diff --git a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py
index 20f4db165..578946cb0 100644
--- a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py
+++ b/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)
diff --git a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py
index a76a5434c..697f0b968 100644
--- a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py
+++ b/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,
diff --git a/vllm_omni/diffusion/diffusion_kv/request.py b/vllm_omni/diffusion/diffusion_kv/request.py
index 83f9bffdb..694af587c 100644
--- a/vllm_omni/diffusion/diffusion_kv/request.py
+++ b/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
diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py
index f15db04bd..9b0f71e23 100644
--- a/vllm_omni/diffusion/executor/abstract.py
+++ b/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."""
diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py
index 1459da495..496f384d7 100644
--- a/vllm_omni/diffusion/forward_context.py
+++ b/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.
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
index 1e91f87c1..743b8af91 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
+++ b/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)",
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
index 441a9636e..4856fac75 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
+++ b/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,
)
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/request_layout.py b/vllm_omni/diffusion/models/hunyuan_image3/request_layout.py
index 32353ed62..6eb873219 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/request_layout.py
+++ b/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)
+ )
diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py
index 9e9a4bfbb..b5be82567 100644
--- a/vllm_omni/diffusion/request.py
+++ b/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
diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py
index d9bb78ea4..4fd695dc5 100644
--- a/vllm_omni/diffusion/sched/base_scheduler.py
+++ b/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.
diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py
index 122a97ec2..fd2b47c63 100644
--- a/vllm_omni/diffusion/sched/interface.py
+++ b/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
diff --git a/vllm_omni/diffusion/sched/request_scheduler.py b/vllm_omni/diffusion/sched/request_scheduler.py
index 4168ae974..826a63c67 100644
--- a/vllm_omni/diffusion/sched/request_scheduler.py
+++ b/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] = {}
diff --git a/vllm_omni/diffusion/sched/step_scheduler.py b/vllm_omni/diffusion/sched/step_scheduler.py
index a5a7a04ae..636c71ba2 100644
--- a/vllm_omni/diffusion/sched/step_scheduler.py
+++ b/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] = {}
diff --git a/vllm_omni/diffusion/vllm_config.py b/vllm_omni/diffusion/vllm_config.py
index c90671d60..12ae1bdb9 100644
--- a/vllm_omni/diffusion/vllm_config.py
+++ b/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)
diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py
index 5fb124a18..6a5b81c58 100644
--- a/vllm_omni/diffusion/worker/diffusion_model_runner.py
+++ b/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)
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index f934646dc..4eadf33db 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/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)
diff --git a/vllm_omni/engine/errors.py b/vllm_omni/engine/errors.py
new file mode 100644
index 000000000..0128176d6
--- /dev/null
+++ b/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."""
diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py
index 411c8b5c2..277a48119 100644
--- a/vllm_omni/engine/orchestrator.py
+++ b/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()
diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py
index f15e1238f..faaf6b707 100644
--- a/vllm_omni/engine/stage_init_utils.py
+++ b/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
diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py
index ce958ba6d..ed22bfa87 100644
--- a/vllm_omni/engine/stage_runtime.py
+++ b/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,
)
diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
index 31225d7ad..0e48412d5 100755
--- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
+++ b/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("![]()
")
self._mrope_boi_token_id = tokenizer.convert_tokens_to_ids("")
self._mrope_eoi_token_id = tokenizer.convert_tokens_to_ids("")
diff --git a/vllm_omni/worker/mixins.py b/vllm_omni/worker/mixins.py
index 5f0fd8300..724cc87be 100644
--- a/vllm_omni/worker/mixins.py
+++ b/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``