[Diffusion] Add native KV cache initialization and Scheduler-managed block allocation (#6094)

Signed-off-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>

docs/configuration/stage_configs.md

@@ -38,11 +38,12 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne
 |-------|------|----------|---------|-------------|
 | `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
 | `max_num_seqs` | int | optional | `64` | Max concurrent sequences per stage. |
-| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage memory budget. |
+| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage total memory target; used for automatic KV-cache sizing. |
+| `kv_cache_memory_bytes` | int \| null | optional | `null` | Explicit per-rank KV-cache budget in bytes; overrides automatic sizing. |
 | `tensor_parallel_size` | int | optional | `1` | TP degree for this stage. |
 | `enforce_eager` | bool | optional | `false` | Disable CUDA graphs. |
-| `max_num_batched_tokens` | int | optional | `32768` | Prefill budget. |
-| `max_model_len` | int \| null | optional | `null` | Per-stage context length (auto-sets `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1` when larger than HF default). |
+| `max_num_batched_tokens` | int | optional | `32768` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. |
+| `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. |
 | `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
 | `devices` | str | optional | `"0"` | `CUDA_VISIBLE_DEVICES`-style device list. |
 | `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |

tests/diffusion/diffusion_kv/test_initialization.py

@@ -0,0 +1,297 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+import vllm_omni.diffusion.diffusion_kv.initialization as initialization
+import vllm_omni.diffusion.vllm_config as diffusion_vllm_config
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
+from vllm_omni.quantization import build_quant_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _Executor:
+    def __init__(self) -> None:
+        self.specs = [{"layer0": object()}, {"layer0": object()}]
+        self.memory = [1024, 1024]
+        self.configs = None
+        self.profile_requests = None
+
+    def get_kv_cache_specs(self):
+        return self.specs
+
+    def determine_available_kv_memory(self, profile_requests):
+        self.profile_requests = profile_requests
+        return self.memory
+
+    def set_kv_cache_configs(self, configs) -> None:
+        self.configs = configs
+
+
+def _od_config(**overrides):
+    values = dict(
+        model="test-diffusion",
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        dtype=torch.bfloat16,
+        max_model_len=64,
+        quantization_config=None,
+        tf_model_config=None,
+        enforce_eager=True,
+        is_moe=False,
+        additional_config={},
+        profiler_config=None,
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+            data_parallel_size=1,
+            enable_expert_parallel=False,
+        ),
+        gpu_memory_utilization=0.9,
+        kv_cache_memory_bytes=None,
+        max_num_batched_tokens=64,
+        num_gpus=1,
+    )
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
+    executor = _Executor()
+    od_config = _od_config(num_gpus=2)
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=4096),
+        max_in_flight_tokens=128,
+    )
+    worker_configs = [object(), object()]
+    scheduler_kv_cache_config = object()
+    profile_requests = [object()]
+    monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: vllm_config)
+    build = Mock(return_value=(worker_configs, scheduler_kv_cache_config))
+    monkeypatch.setattr(
+        initialization,
+        "build_native_kv_cache_configs",
+        build,
+    )
+    monkeypatch.setattr(initialization, "resolve_kv_cache_block_sizes", lambda *args: (16, 16))
+
+    result = initialization.initialize_diffusion_kv_control_plane(
+        executor,
+        od_config,
+        profile_requests=profile_requests,
+        device=torch.device("cpu"),
+    )
+
+    build.assert_called_once_with(vllm_config, executor.specs, executor.memory)
+    assert result == (scheduler_kv_cache_config, 16, 16, vllm_config)
+    assert executor.configs == worker_configs
+    assert executor.profile_requests is profile_requests
+
+
+def test_dense_mode_skips_worker_kv_initialization(monkeypatch) -> None:
+    executor = _Executor()
+    od_config = _od_config(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
+    get_device = Mock(side_effect=AssertionError("dense mode must not resolve a KV device"))
+    monkeypatch.setattr(initialization.current_omni_platform, "get_torch_device", get_device)
+
+    result = initialization.initialize_diffusion_kv_control_plane(executor, od_config)
+
+    assert result is None
+    assert executor.configs is None
+    get_device.assert_not_called()
+
+
+def test_control_plane_rejects_rank_count_mismatch(monkeypatch) -> None:
+    executor = _Executor()
+    executor.memory = [1024]
+    od_config = _od_config(num_gpus=2)
+    monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: object())
+
+    with pytest.raises(ValueError, match="rank count mismatch"):
+        initialization.initialize_diffusion_kv_control_plane(
+            executor,
+            od_config,
+            profile_requests=[object()],
+            device=torch.device("cpu"),
+        )
+
+
+def test_paged_control_plane_requires_prepared_profile_request() -> None:
+    with pytest.raises(ValueError, match="prepared memory profile requests"):
+        initialization.initialize_diffusion_kv_control_plane(
+            _Executor(),
+            _od_config(),
+            device=torch.device("cpu"),
+        )
+
+
+def test_native_config_builder_produces_matching_worker_and_scheduler_pools() -> None:
+    od_config = _od_config()
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
+    spec = FullAttentionSpec(
+        block_size=vllm_config.cache_config.block_size,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+        non_causal=True,
+    )
+    num_blocks = 32
+
+    worker_configs, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
+        vllm_config,
+        [{"layer0": spec}],
+        [spec.page_size_bytes * num_blocks],
+    )
+
+    assert len(worker_configs) == 1
+    assert worker_configs[0].num_blocks == num_blocks
+    assert scheduler_kv_cache_config.num_blocks == num_blocks
+    assert worker_configs[0].kv_cache_groups == scheduler_kv_cache_config.kv_cache_groups
+
+
+def test_native_control_plane_config_initializes_scheduler() -> None:
+    od_config = _od_config()
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
+    spec = FullAttentionSpec(
+        block_size=vllm_config.cache_config.block_size,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+        non_causal=True,
+    )
+    _, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
+        vllm_config,
+        [{"layer0": spec}],
+        [spec.page_size_bytes * 32],
+    )
+
+    scheduler = RequestScheduler()
+    scheduler.initialize(
+        od_config,
+        kv_cache_config=scheduler_kv_cache_config,
+        scheduler_block_size=spec.block_size,
+        hash_block_size=spec.block_size,
+        kv_vllm_config=vllm_config,
+    )
+
+    assert scheduler._diffusion_kv_manager is not None
+    scheduler.close()
+
+
+def test_dense_config_keeps_new_paged_sizing_inputs_inactive() -> None:
+    od_config = _od_config(
+        diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=1,
+            pipeline_parallel_size=4,
+            data_parallel_size=1,
+            enable_expert_parallel=False,
+        ),
+        gpu_memory_utilization=0.5,
+        kv_cache_memory_bytes=4096,
+        max_num_batched_tokens=64,
+    )
+    vllm_config = diffusion_vllm_config.create_base_diffusion_vllm_config(torch.device("cpu"), od_config)
+    original_pp_size = vllm_config.parallel_config.pipeline_parallel_size
+    original_memory_utilization = vllm_config.cache_config.gpu_memory_utilization
+    original_memory_bytes = vllm_config.cache_config.kv_cache_memory_bytes
+    original_token_budget = vllm_config.scheduler_config.max_num_batched_tokens
+
+    diffusion_vllm_config.configure_diffusion_vllm_config(vllm_config, od_config)
+
+    assert vllm_config.parallel_config.pipeline_parallel_size == original_pp_size
+    assert vllm_config.cache_config.gpu_memory_utilization == original_memory_utilization
+    assert vllm_config.cache_config.kv_cache_memory_bytes == original_memory_bytes
+    assert vllm_config.scheduler_config.max_num_batched_tokens == original_token_budget
+
+
+def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -> None:
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
+        torch.device("cpu"),
+        _od_config(gpu_memory_utilization=0.42),
+    )
+
+    assert vllm_config.cache_config.gpu_memory_utilization == 0.42
+
+
+def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None:
+    quantization_config = build_quant_config(
+        {
+            "quant_method": "modelopt",
+            "quant_algo": "FP8",
+            "ignore": [],
+        }
+    )
+    model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+        _od_config(
+            diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+            quantization_config=quantization_config,
+        )
+    )
+
+    assert model_config.dtype is torch.bfloat16
+    assert model_config.quantization == "modelopt"
+    assert model_config.quantization_config is quantization_config
+    assert model_config.is_quantized
+    assert model_config.is_moe is False
+    assert model_config.is_nvfp4_quantized() is False
+    assert model_config.original_max_model_len == model_config.max_model_len
+
+    unquantized_model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+        _od_config(
+            diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+            quantization_config=None,
+        )
+    )
+    assert not unquantized_model_config.is_quantized
+
+
+def test_minus_one_max_model_len_uses_model_limit() -> None:
+    od_config = SimpleNamespace(
+        max_model_len=-1,
+        tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+    )
+
+    assert diffusion_vllm_config.resolve_diffusion_max_model_len(od_config) == 4096
+
+    model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+        _od_config(
+            max_model_len=-1,
+            tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+        )
+    )
+    assert model_config.max_model_len == 4096
+    assert model_config.original_max_model_len == -1
+
+
+def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None:
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
+        torch.device("cpu"),
+        _od_config(
+            max_model_len=-1,
+            tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+        ),
+    )
+    spec = FullAttentionSpec(
+        block_size=vllm_config.cache_config.block_size,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+        non_causal=True,
+    )
+
+    initialization.build_native_kv_cache_configs(
+        vllm_config,
+        [{"layer0": spec}],
+        [spec.page_size_bytes * 16],
+    )
+
+    assert vllm_config.model_config.original_max_model_len == -1
+    assert vllm_config.model_config.max_model_len == 256

tests/diffusion/diffusion_kv/test_manager.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+
+from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+BLOCK_SIZE = 4
+
+
+def _config(num_blocks: int) -> KVCacheConfig:
+    native_kv_managers.register_all_kvcache_specs(None)
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+    )
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
+    )
+
+
+def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest:
+    return DiffusionKVRequest(
+        f"{public_id}/diffusion-kv/{sequence_id}",
+        sequence_id=sequence_id,
+        prefix_len=4,
+        target_len=4,
+        seq_len=seq_len,
+        kv_contexts=kv_contexts,
+    )
+
+
+def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager:
+    return DiffusionKVCacheManager(
+        _config(num_blocks),
+        max_model_len=max_model_len,
+        scheduler_block_size=BLOCK_SIZE,
+        hash_block_size=BLOCK_SIZE,
+    )
+
+
+def test_reserve_and_free_multi_cfg_request() -> None:
+    manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    requests = (_request("public", 0), _request("public", 1))
+
+    metadata = manager.reserve_request("public", requests)
+
+    assert metadata is not None
+    assert metadata.request_id == "public"
+    assert metadata.allocation_generation == 1
+    assert [sequence.sequence_id for sequence in metadata.sequences] == [0, 1]
+    assert [sequence.prefix_len for sequence in metadata.sequences] == [4, 4]
+    assert [sequence.target_len for sequence in metadata.sequences] == [4, 4]
+    assert [sequence.seq_len for sequence in metadata.sequences] == [8, 8]
+    assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
+    assert all(sequence.context_ids == () for sequence in metadata.sequences)
+    assert metadata.contexts == ()
+    assert metadata == manager.get_metadata("public")
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before - 4
+
+    manager.free_request("public")
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_impossible_cfg_allocation_rolls_back_and_fails_fast() -> None:
+    manager = _manager(3)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+
+    with pytest.raises(DiffusionKVAdmissionError, match="cannot fit even when the block pool is empty"):
+        manager.reserve_request(
+            "public",
+            (_request("public", 0), _request("public", 1)),
+        )
+
+    assert manager.has_request("public") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_impossible_cfg_allocation_fails_fast_while_pool_is_busy() -> None:
+    manager = _manager(5)
+    assert manager.reserve_request("running", (_request("running", 0),)) is not None
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+
+    with pytest.raises(DiffusionKVAdmissionError, match="required_blocks=6, available_blocks=4"):
+        manager.reserve_request(
+            "impossible",
+            tuple(_request("impossible", sequence_id) for sequence_id in range(3)),
+        )
+
+    assert manager.has_request("impossible") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_temporary_capacity_pressure_returns_none() -> None:
+    manager = _manager(3)
+    assert manager.reserve_request("running", (_request("running", 0),)) is not None
+
+    allocation = manager.reserve_request("waiting", (_request("waiting", 0),))
+
+    assert allocation is None
+    assert manager.has_request("waiting") is False
+
+
+def test_partial_cfg_rollback_can_retry_after_capacity_is_released() -> None:
+    manager = _manager(5)
+    assert manager.reserve_request("running", (_request("running", 0),)) is not None
+    waiting = (_request("waiting", 0), _request("waiting", 1))
+
+    assert manager.reserve_request("waiting", waiting) is None
+    assert manager.has_request("waiting") is False
+
+    manager.free_request("running")
+    metadata = manager.reserve_request("waiting", waiting)
+
+    assert metadata is not None
+    assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
+
+
+def test_reallocated_request_gets_a_new_allocation_generation() -> None:
+    manager = _manager(3)
+    request = (_request("public", 0),)
+
+    first = manager.reserve_request("public", request)
+    assert first is not None
+    manager.free_request("public")
+    second = manager.reserve_request("public", request)
+
+    assert second is not None
+    assert second.allocation_generation == first.allocation_generation + 1
+
+
+def test_cfg_allocation_exception_rolls_back(monkeypatch) -> None:
+    manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    native_allocate = manager.native_manager.allocate_slots
+    num_calls = 0
+
+    def fail_second_allocation(*args, **kwargs):
+        nonlocal num_calls
+        num_calls += 1
+        if num_calls == 2:
+            raise RuntimeError("injected allocation failure")
+        return native_allocate(*args, **kwargs)
+
+    monkeypatch.setattr(manager.native_manager, "allocate_slots", fail_second_allocation)
+
+    with pytest.raises(RuntimeError, match="injected allocation failure"):
+        manager.reserve_request(
+            "public",
+            (_request("public", 0), _request("public", 1)),
+        )
+
+    assert manager.has_request("public") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_rejects_independent_context_until_role_routing_is_implemented() -> None:
+    manager = _manager(8)
+    context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
+
+    with pytest.raises(DiffusionKVAdmissionError, match="DiffusionKVContext"):
+        manager.reserve_request(
+            "public",
+            (_request("public", 0, kv_contexts=(context,)),),
+        )
+
+
+def test_rejects_sequence_longer_than_admission_bound() -> None:
+    manager = _manager(8, max_model_len=8)
+
+    with pytest.raises(DiffusionKVAdmissionError, match="exceeds max_model_len"):
+        manager.reserve_request("public", (_request("public", 0, seq_len=9),))
+
+
+@pytest.mark.parametrize(
+    "requests",
+    [
+        (_request("public", 1),),
+        (_request("public", 1), _request("public", 0)),
+        (_request("public", 0), _request("public", 2)),
+    ],
+)
+def test_rejects_ambiguous_sequence_order(requests) -> None:
+    manager = _manager(8)
+
+    with pytest.raises(ValueError, match="ordered by contiguous sequence_id"):
+        manager.reserve_request("public", requests)
+
+
+def test_close_releases_every_public_request() -> None:
+    manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    assert manager.reserve_request("first", (_request("first", 0),)) is not None
+    assert manager.reserve_request("second", (_request("second", 0),)) is not None
+
+    manager.close()
+
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+    assert manager.has_request("first") is False
+    assert manager.has_request("second") is False

tests/diffusion/diffusion_kv/test_request.py

@@ -107,11 +107,8 @@ def test_request_rejects_non_context_values() -> None:
 def _manager(*, enable_caching: bool) -> KVCacheManager:
     # vLLM 0.26 lazily registers its built-in KV specs. Another test may have
     # already registered an out-of-tree spec, making the registry non-empty
-    # before FullAttentionSpec is registered. Initialize the native mappings
-    # explicitly; vLLM 0.25 has no registration hook and needs no setup.
-    register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None)
-    if register_all is not None:
-        register_all(None)
+    # before FullAttentionSpec is registered. Initialize the native mappings.
+    native_kv_managers.register_all_kvcache_specs(None)
 
     spec = FullAttentionSpec(
         block_size=BLOCK_SIZE,

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from torch import nn
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+
+import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _Backend:
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        return True
+
+
+def _attention(*, enabled: bool) -> Attention:
+    attention = Attention.__new__(Attention)
+    nn.Module.__init__(attention)
+    attention.paged_kv_cache_role = "primary" if enabled else None
+    attention.paged_kv_cache_dtype = torch.bfloat16
+    attention.num_kv_heads = 2
+    attention.head_size = 8
+    attention.causal = False
+    attention.attn_backend = _Backend
+    return attention
+
+
+def _runner(attention: Attention) -> DiffusionModelRunner:
+    runner = object.__new__(DiffusionModelRunner)
+    runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
+    runner.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=4),
+        model_config=SimpleNamespace(dtype=torch.float16),
+    )
+    pipeline = nn.Module()
+    pipeline.image_attention = attention
+    runner.pipeline = pipeline
+    return runner
+
+
+def test_runner_discovers_native_spec_from_loaded_attention() -> None:
+    specs = _runner(_attention(enabled=True)).get_kv_cache_spec()
+
+    assert set(specs) == {"image_attention"}
+    spec = specs["image_attention"]
+    assert isinstance(spec, FullAttentionSpec)
+    assert spec.block_size == 4
+    assert spec.num_kv_heads == 2
+    assert spec.head_size == 8
+    assert spec.dtype is torch.bfloat16
+    assert spec.indexes_kv_by_block_stride is True
+    assert spec.non_causal is True
+
+
+def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None:
+    runner = _runner(_attention(enabled=False))
+
+    with pytest.raises(RuntimeError, match="no cache-enabled Attention"):
+        runner.get_kv_cache_spec()
+
+
+def test_runner_retains_matching_rank_local_config() -> None:
+    runner = _runner(_attention(enabled=True))
+    spec = runner.get_kv_cache_spec()["image_attention"]
+    config = KVCacheConfig(
+        num_blocks=8,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)],
+    )
+
+    runner.set_kv_cache_config(config)
+
+    assert runner.kv_cache_config is config
+
+
+def test_runner_rejects_rank_local_config_for_different_layers() -> None:
+    runner = _runner(_attention(enabled=True))
+    spec = runner.get_kv_cache_spec()["image_attention"]
+    config = KVCacheConfig(
+        num_blocks=8,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)],
+    )
+
+    with pytest.raises(ValueError, match="layer mismatch"):
+        runner.set_kv_cache_config(config)
+
+
+def test_worker_selects_its_rank_local_config() -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 1
+    worker.od_config = SimpleNamespace(num_gpus=2)
+    worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config))
+    configs = [object(), object()]
+
+    worker.set_kv_cache_configs(configs)
+
+    assert worker.installed is configs[1]
+    with pytest.raises(ValueError, match="rank count mismatch"):
+        worker.set_kv_cache_configs(configs[:1])
+
+
+def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=4096))
+    worker.init_snapshot = SimpleNamespace(free_memory=16384)
+    worker.requested_memory = 8192
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=1024)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_request = object()
+
+    assert worker.determine_available_kv_memory(profile_request) == [4096]
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.init_snapshot = SimpleNamespace(free_memory=900)
+    worker.requested_memory = 750
+    worker.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            kv_cache_memory_bytes=0,
+            gpu_memory_utilization=0.75,
+        )
+    )
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_result = SimpleNamespace(
+        non_kv_cache_memory=400,
+        after_profile=SimpleNamespace(free_memory=500),
+    )
+    monkeypatch.setattr(
+        diffusion_worker_module,
+        "memory_profiling",
+        lambda snapshot, weights_memory: nullcontext(profile_result),
+    )
+    profile_request = object()
+
+    assert worker.determine_available_kv_memory(profile_request) == [350]
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+@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:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.init_snapshot = SimpleNamespace(free_memory=900)
+    worker.requested_memory = 750
+    worker.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            kv_cache_memory_bytes=None,
+            gpu_memory_utilization=0.75,
+        )
+    )
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_result = SimpleNamespace(
+        non_kv_cache_memory=non_kv_cache_memory,
+        after_profile=SimpleNamespace(free_memory=100),
+    )
+    monkeypatch.setattr(
+        diffusion_worker_module,
+        "memory_profiling",
+        lambda snapshot, weights_memory: nullcontext(profile_result),
+    )
+    profile_request = object()
+
+    with pytest.raises(
+        RuntimeError,
+        match=(
+            "rank 0: RuntimeError: No memory remains for Diffusion KV cache after profiling: "
+            f"requested_memory=750 bytes, non_kv_cache_memory={non_kv_cache_memory} bytes"
+        ),
+    ):
+        worker.determine_available_kv_memory(profile_request)
+
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_profiles_activation_headroom_instead_of_current_residency(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.init_snapshot = SimpleNamespace(free_memory=1200)
+    worker.requested_memory = 1000
+    worker.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            kv_cache_memory_bytes=None,
+            gpu_memory_utilization=0.8,
+        )
+    )
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=250)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_result = SimpleNamespace(
+        non_kv_cache_memory=650,
+        after_profile=SimpleNamespace(free_memory=550),
+    )
+    profiling_calls = []
+
+    def profile(snapshot, *, weights_memory):
+        profiling_calls.append((snapshot, weights_memory))
+        return nullcontext(profile_result)
+
+    monkeypatch.setattr(
+        diffusion_worker_module,
+        "memory_profiling",
+        profile,
+    )
+    profile_request = object()
+
+    assert worker.determine_available_kv_memory(profile_request) == [350]
+    assert profiling_calls == [(worker.init_snapshot, 250)]
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_rejects_profile_without_preload_snapshot(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.init_snapshot = None
+    worker.requested_memory = None
+    worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=None))
+    worker.model_runner = SimpleNamespace()
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+
+    with pytest.raises(RuntimeError, match="snapshot was not captured before model loading"):
+        worker.determine_available_kv_memory(object())
+
+
+def test_rank_probe_gathers_local_failure_before_raising(monkeypatch) -> None:
+    gathered = []
+
+    def gather(local_result):
+        gathered.append(local_result)
+        return [local_result, (True, "peer-result")]
+
+    def fail_local_probe():
+        raise ValueError("local failure")
+
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", gather)
+
+    with pytest.raises(RuntimeError, match="rank 0: ValueError: local failure"):
+        diffusion_worker_module._run_and_gather_rank_values(
+            "test probe",
+            fail_local_probe,
+        )
+
+    assert gathered == [(False, "ValueError: local failure")]

tests/diffusion/diffusion_kv/test_worker_contract.py

@@ -85,19 +85,15 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None:
     assert new_req.diffusion_kv_metadata is metadata
 
 
-def test_paged_request_without_metadata_reaches_request_forward() -> None:
+def test_paged_request_without_metadata_fails_before_request_forward() -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
-    expected_output = DiffusionOutput(output=None)
-    runner._execute_request_list = Mock(
-        return_value=SimpleNamespace(
-            runner_outputs=[SimpleNamespace(result=expected_output)],
-        )
-    )
+    runner._execute_request_list = Mock()
     req = SimpleNamespace(request_id="req-0")
 
-    assert runner.execute_model(req) is expected_output
+    with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
+        runner.execute_model(req)
 
-    runner._execute_request_list.assert_called_once()
+    runner._execute_request_list.assert_not_called()
 
 
 def test_dense_request_rejects_metadata_before_request_forward() -> None:
@@ -141,18 +137,18 @@ def test_request_rpc_rejects_envelope_request_identity_mismatch() -> None:
     assert calls == []
 
 
-def test_batch_path_allows_missing_metadata_before_forward() -> None:
+def test_batch_path_rejects_missing_metadata_before_forward() -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
-    expected_output = object()
-    runner._execute_request_list = Mock(return_value=expected_output)
+    runner._execute_request_list = Mock()
     new_req = NewRequestData(
         request_id="req-0",
         req=SimpleNamespace(request_id="req-0"),
     )
 
-    assert runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config) is expected_output
+    with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
+        runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config)
 
-    runner._execute_request_list.assert_called_once()
+    runner._execute_request_list.assert_not_called()
 
 
 def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
@@ -170,7 +166,7 @@ def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
     runner._execute_request_list.assert_not_called()
 
 
-def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
+def test_step_path_rejects_missing_metadata_before_step_support_check() -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
     runner.pipeline = object()
     runner._supports_step_mode = Mock(return_value=False)
@@ -179,10 +175,10 @@ def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
         req=SimpleNamespace(request_id="req-0"),
     )
 
-    with pytest.raises(ValueError, match="does not support step execution"):
+    with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
         runner.execute_stepwise(make_scheduler_output(new_req))
 
-    runner._supports_step_mode.assert_called_once()
+    runner._supports_step_mode.assert_not_called()
 
 
 def test_step_path_rejects_envelope_request_identity_mismatch() -> None:
@@ -288,3 +284,42 @@ def test_dense_worker_call_does_not_extend_model_runner_signature() -> None:
     worker.execute_model(req, SimpleNamespace())
 
     assert calls == [(req, None)]
+
+
+def test_dlo_worker_selects_request_and_metadata_from_same_envelope(monkeypatch) -> None:
+    calls: list[tuple] = []
+
+    class ModelRunner:
+        def execute_model(self, req, **kwargs):
+            calls.append((req, kwargs))
+            return DiffusionOutput(output=None)
+
+    worker = object.__new__(DiffusionWorker)
+    worker.model_runner = ModelRunner()
+    worker.lora_manager = None
+    worker._get_profiler = lambda: None
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.distributed.parallel_state.get_data_parallel_rank",
+        lambda: 1,
+    )
+    req_0 = SimpleNamespace(request_id="req-0")
+    req_1 = SimpleNamespace(request_id="req-1")
+    metadata_0 = make_metadata("req-0")
+    metadata_1 = make_metadata("req-1")
+    envelopes = [
+        NewRequestData(request_id="req-0", req=req_0, diffusion_kv_metadata=metadata_0),
+        NewRequestData(request_id="req-1", req=req_1, diffusion_kv_metadata=metadata_1),
+    ]
+
+    result = worker.execute_model(envelopes, SimpleNamespace())
+
+    assert calls == [
+        (
+            req_1,
+            {
+                "kv_prefetch_job": None,
+                "diffusion_kv_metadata": metadata_1,
+            },
+        )
+    ]
+    assert result["dp_rank"] == 1

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -118,6 +118,14 @@ def _call_mgr(
     )
 
 
+def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
+    mgr = _make_cache_mgr()
+
+    assert isinstance(mgr, nn.Module)
+    assert dict(mgr.named_modules())["attn"] is mgr.attn
+    assert mgr.state_dict() == {}
+
+
 # ============================================================
 # Test 1: No AR KV — basic cache → reuse
 # ============================================================

tests/diffusion/test_diffusion_config_propagation.py

@@ -130,6 +130,21 @@ class TestCreateDefaultDiffusion:
 
         assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
 
+    def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch):
+        monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500)
+        od = _roundtrip_diffusion_config(
+            model="x",
+            kv_cache_memory_bytes=4096,
+            gpu_memory_utilization=0.75,
+            max_num_batched_tokens=2048,
+            max_model_len=4096,
+        )
+
+        assert od.kv_cache_memory_bytes == 4096
+        assert od.gpu_memory_utilization == 0.75
+        assert od.max_num_batched_tokens == 2048
+        assert od.max_model_len == 4096
+
 
 def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
     od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -161,9 +161,83 @@ def test_scheduler_initialization_failure_closes_scheduler_and_executor(monkeypa
     assert exc_info.value is initialization_error
     custom_scheduler.initialize.assert_called_once_with(od_config)
     custom_scheduler.close.assert_called_once_with()
+
+
+def test_init_shuts_down_executor_when_kv_control_plane_initialization_fails(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    fake_executor = SimpleNamespace(shutdown=Mock())
+    monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_resolve_execution_mode",
+        lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
+    )
+    monkeypatch.setattr(
+        DiffusionEngine, "_init_executor", lambda self, config: setattr(self, "executor", fake_executor)
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+        Mock(side_effect=RuntimeError("KV initialization failed")),
+    )
+
+    with pytest.raises(RuntimeError, match="KV initialization failed"):
+        DiffusionEngine(SimpleNamespace())
+
     fake_executor.shutdown.assert_called_once_with()
 
 
+def test_paged_init_profiles_before_scheduler_initialization(monkeypatch: pytest.MonkeyPatch) -> None:
+    events = []
+    expected_profile_requests = [object()]
+    fake_executor = SimpleNamespace(shutdown=Mock())
+    kv_cache_config = object()
+    kv_vllm_config = object()
+    od_config = SimpleNamespace(diffusion_kv_mode="paged_scheduler")
+
+    monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_resolve_execution_mode",
+        lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
+    )
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_init_executor",
+        lambda self, config: setattr(self, "executor", fake_executor),
+    )
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_prepare_diffusion_kv_profile_requests",
+        lambda self: events.append("prepare-profile") or expected_profile_requests,
+    )
+
+    def initialize(executor, config, *, profile_requests):
+        events.append("profile-workers")
+        assert executor is fake_executor
+        assert config is od_config
+        assert profile_requests is expected_profile_requests
+        return kv_cache_config, 16, 16, kv_vllm_config
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+        initialize,
+    )
+
+    def init_scheduler(self, config, scheduler, *args, **kwargs):
+        events.append("initialize-scheduler")
+        self.scheduler = SimpleNamespace(close=Mock())
+
+    monkeypatch.setattr(DiffusionEngine, "_init_scheduler", init_scheduler)
+    monkeypatch.setattr(DiffusionEngine, "_init_runtime_state", lambda self: None)
+    monkeypatch.setattr(DiffusionEngine, "_init_execute_fn", lambda self: None)
+    monkeypatch.setattr(DiffusionEngine, "_log_execution_mode", lambda self, config: None)
+
+    DiffusionEngine(od_config)
+
+    assert events == ["prepare-profile", "profile-workers", "initialize-scheduler"]
+
+
 @pytest.mark.asyncio
 async def test_step_compatibility_wrapper_returns_final_batch() -> None:
     engine = _make_engine()

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -1,9 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+from unittest.mock import Mock
+
 import pytest
 
 from vllm_omni.diffusion import io_support
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -55,3 +63,151 @@ def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.Monk
     )
 
     assert io_support.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2
+
+
+def test_dummy_run_image_count_resolves_hunyuan_architecture_alias() -> None:
+    assert io_support.get_dummy_run_num_image_inputs("HunyuanImage3ForCausalMM") == 3
+    assert io_support.get_dummy_run_num_image_inputs("unknown") == 1
+
+
+def test_dense_mode_does_not_build_kv_profile_request() -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
+    engine._make_dummy_request = Mock(side_effect=AssertionError("dense mode must not profile"))
+
+    assert engine._prepare_diffusion_kv_profile_requests() is None
+    engine._make_dummy_request.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("execution_mode", "uses_dlo_dp", "max_num_seqs", "expected_profile_requests"),
+    [
+        (DiffusionExecutionMode.STEP_BATCH, False, 2, 2),
+        (DiffusionExecutionMode.REQUEST_BATCH, True, 4, 1),
+        (DiffusionExecutionMode.STEP_BATCH, True, 3, 3),
+    ],
+)
+def test_paged_kv_profile_requests_match_per_rank_batch(
+    execution_mode: DiffusionExecutionMode,
+    uses_dlo_dp: bool,
+    max_num_seqs: int,
+    expected_profile_requests: int,
+) -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine.execution_mode = execution_mode
+    engine.od_config = SimpleNamespace(
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        model_class_name="HunyuanImage3ForCausalMM",
+        max_num_seqs=max_num_seqs,
+        parallel_config=SimpleNamespace(data_parallel_size=max_num_seqs if uses_dlo_dp else 1),
+        enable_distributed_layerwise_offload=uses_dlo_dp,
+        dlo_use_allgather=True,
+    )
+    request = OmniDiffusionRequest(
+        prompt="profile",
+        request_id="profile-request",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+    )
+    engine._make_dummy_request = Mock(return_value=request)
+    prepared_layout = object()
+    scheduler_kv_state = (
+        DiffusionKVRequest(
+            "profile-request/diffusion-kv/0",
+            sequence_id=0,
+            prefix_len=3,
+            target_len=8,
+            seq_len=12,
+        ),
+        DiffusionKVRequest(
+            "profile-request/diffusion-kv/1",
+            sequence_id=1,
+            prefix_len=4,
+            target_len=8,
+            seq_len=13,
+        ),
+    )
+
+    def preprocess(req: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        req.prepared_layout = prepared_layout
+        req.diffusion_kv_requests = scheduler_kv_state
+        return req
+
+    engine._prepare_request_for_admission = Mock(side_effect=preprocess)
+
+    result = engine._prepare_diffusion_kv_profile_requests()
+
+    assert result is not None
+    assert len(result) == expected_profile_requests
+    assert len({profile_request.request_id for profile_request in result}) == expected_profile_requests
+    assert all(profile_request.is_dummy_run() for profile_request in result)
+    assert all(profile_request.prepared_layout is prepared_layout for profile_request in result)
+    assert all(profile_request.diffusion_kv_requests is None for profile_request in result)
+    assert len({id(profile_request.sampling_params) for profile_request in result}) == expected_profile_requests
+    assert engine._diffusion_kv_profile_limits == (2, 13, 8)
+    assert len(result) * engine._diffusion_kv_profile_limits[0] == expected_profile_requests * 2
+    engine._make_dummy_request.assert_called_once_with(
+        height=1024,
+        width=1024,
+        guidance_scale=5.0,
+        num_image_inputs=3,
+    )
+    engine._prepare_request_for_admission.assert_called_once_with(request)
+
+
+@pytest.mark.parametrize(
+    ("num_sequences", "seq_len", "target_len"),
+    [
+        (3, 16, 8),
+        (1, 17, 8),
+        (1, 16, 9),
+    ],
+)
+def test_paged_kv_admission_rejects_shape_beyond_profile_envelope(
+    num_sequences: int,
+    seq_len: int,
+    target_len: int,
+) -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine._diffusion_kv_profile_limits = (2, 16, 8)
+    engine.pre_process_func = lambda request: request
+    request = OmniDiffusionRequest(
+        prompt="too large",
+        request_id="too-large",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        diffusion_kv_requests=tuple(
+            DiffusionKVRequest(
+                f"too-large/diffusion-kv/{sequence_id}",
+                sequence_id=sequence_id,
+                prefix_len=0,
+                target_len=target_len,
+                seq_len=seq_len,
+            )
+            for sequence_id in range(num_sequences)
+        ),
+    )
+
+    with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
+        engine._prepare_request_for_admission(request)
+
+
+def test_paged_kv_admission_accepts_shape_at_profile_envelope() -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine._diffusion_kv_profile_limits = (2, 16, 8)
+    engine.pre_process_func = lambda request: request
+    request = OmniDiffusionRequest(
+        prompt="fits",
+        request_id="fits",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        diffusion_kv_requests=tuple(
+            DiffusionKVRequest(
+                f"fits/diffusion-kv/{sequence_id}",
+                sequence_id=sequence_id,
+                prefix_len=8,
+                target_len=8,
+                seq_len=16,
+            )
+            for sequence_id in range(2)
+        ),
+    )
+
+    assert engine._prepare_request_for_admission(request) is request

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -473,6 +473,38 @@ async def test_cancelled_future_does_not_kill_busy_loop():
         _stop_engine(engine)
 
 
+@pytest.mark.asyncio
+async def test_native_kv_reservation_error_wakes_stream_without_killing_busy_loop():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+
+    def fail_reservation(*args, **kwargs):
+        raise RuntimeError("native allocation bug")
+
+    engine.scheduler._diffusion_kv_manager = SimpleNamespace(
+        has_request=lambda request_id: False,
+        reserve_request=fail_reservation,
+        free_request=lambda request_id: None,
+    )
+    request = _make_request("kv-error")
+    request.diffusion_kv_requests = (object(),)
+    try:
+        response_stream = engine.async_add_req_and_stream_response(request)
+        output_queue = engine._out_streams[request.request_id]
+        for _ in range(300):
+            if not output_queue.empty():
+                break
+            await asyncio.sleep(0.01)
+        assert not output_queue.empty(), "terminal KV allocation error was not delivered to the request stream"
+
+        output = await anext(response_stream)
+        assert output.error == "native allocation bug"
+        assert output.finished
+        assert engine.worker_thread.is_alive()
+    finally:
+        _stop_engine(engine)
+
+
 # ─────────────────────────── _RpcTask basics ───────────────────────────────
 
 

tests/diffusion/test_diffusion_model_runner.py

@@ -3,6 +3,7 @@
 
 from contextlib import contextmanager
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -512,6 +513,84 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
     assert runner.pipeline.last_req.num_reqs == 1
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatch):
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    request = _make_request()
+    runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
+    record_names = []
+    original_execute = runner._execute_request_list
+
+    def execute(*args, **kwargs):
+        record_names.append(kwargs["record_name"])
+        return original_execute(*args, **kwargs)
+
+    runner._execute_request_list = execute
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    reset_peak_memory_stats = Mock()
+    monkeypatch.setattr(
+        model_runner_module.current_omni_platform,
+        "reset_peak_memory_stats",
+        reset_peak_memory_stats,
+    )
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+    synchronize = Mock()
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
+
+    runner.profile_run([request])
+
+    assert runner.pipeline.forward_calls == 1
+    assert record_names == ["pipeline_memory_profile"]
+    runner._validate_diffusion_kv_metadata.assert_not_called()
+    reset_peak_memory_stats.assert_not_called()
+    synchronize.assert_called_once_with()
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypatch):
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    runner.pipeline = _FinalOnlyStepPipeline()
+    runner.od_config.step_execution = True
+    requests = [_make_request(), _make_request()]
+    requests[0].request_id = "profile-0"
+    requests[1].request_id = "profile-1"
+    for request in requests:
+        request.sampling_params.num_inference_steps = 1
+
+    observed_batch_rows = []
+    original_denoise_step = runner.pipeline.denoise_step
+
+    def denoise_step(input_batch, states):
+        observed_batch_rows.append((len(states), int(input_batch.latents.shape[0])))
+        return original_denoise_step(input_batch, states)
+
+    runner.pipeline.denoise_step = denoise_step
+    runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    reset_peak_memory_stats = Mock()
+    monkeypatch.setattr(
+        model_runner_module.current_omni_platform,
+        "reset_peak_memory_stats",
+        reset_peak_memory_stats,
+    )
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+    synchronize = Mock()
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
+
+    runner.profile_run(requests)
+
+    assert observed_batch_rows == [(2, 2)]
+    assert runner.state_cache == {}
+    assert runner.input_batch is None
+    runner._validate_diffusion_kv_metadata.assert_not_called()
+    reset_peak_memory_stats.assert_not_called()
+    synchronize.assert_called_once_with()
+
+
 class _BatchPipeline:
     """Pipeline returning a configurable list of outputs from forward()."""
 

tests/diffusion/test_diffusion_request.py

@@ -52,6 +52,7 @@ def test_dummy_run_request_is_identified_by_reserved_request_id():
 
     assert req.is_dummy_run()
     assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
+    assert OmniDiffusionRequest.is_dummy_run_request_id(f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-1")
 
 
 def test_non_dummy_request_is_not_identified_as_dummy_run():

tests/diffusion/test_diffusion_scheduler.py

@@ -8,12 +8,14 @@ from types import SimpleNamespace
 
 import pytest
 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, KVCacheTensor
 
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
-from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
     BaseScheduler,
@@ -87,6 +89,52 @@ def _cached_ids(sched_output) -> list[str]:
     return list(sched_output.scheduled_cached_reqs.request_ids)
 
 
+def _initialize_paged_scheduler(
+    scheduler: BaseScheduler,
+    *,
+    num_blocks: int = 64,
+    max_num_seqs: int = 1,
+) -> None:
+    native_kv_managers.register_all_kvcache_specs(None)
+    spec = FullAttentionSpec(
+        block_size=4,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+    )
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
+    )
+    scheduler.initialize(
+        SimpleNamespace(
+            diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+            max_model_len=64,
+            max_num_seqs=max_num_seqs,
+        ),
+        kv_cache_config=config,
+        scheduler_block_size=4,
+        hash_block_size=4,
+        kv_vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(max_model_len=64),
+            max_in_flight_tokens=64,
+        ),
+    )
+
+
+def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) -> None:
+    request.diffusion_kv_requests = (
+        DiffusionKVRequest(
+            f"{request.request_id}/diffusion-kv/0",
+            sequence_id=0,
+            prefix_len=4,
+            target_len=4,
+            seq_len=seq_len,
+        ),
+    )
+
+
 class _StubScheduler:
     def __init__(self, request: OmniDiffusionRequest, output) -> None:
         self._request = request
@@ -400,18 +448,16 @@ class TestRequestScheduler:
         assert self.scheduler.has_requests() is False
 
     def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None:
-        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        _initialize_paged_scheduler(self.scheduler)
         request = _make_request("diffusion-kv")
         prepared_layout = object()
         request.prepared_layout = prepared_layout
-        text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
         kv_request = DiffusionKVRequest(
             "diffusion-kv/diffusion-kv/0",
             sequence_id=0,
             prefix_len=4,
             target_len=8,
             seq_len=16,
-            kv_contexts=(text_context,),
         )
         request.diffusion_kv_requests = (kv_request,)
 
@@ -421,12 +467,179 @@ class TestRequestScheduler:
 
         assert state is not None
         assert state.diffusion_kv_requests == (kv_request,)
-        assert state.diffusion_kv_requests[0].kv_contexts == (text_context,)
         assert request.diffusion_kv_requests is None
         assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout
+        metadata = scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata
+        assert metadata is not None
+        assert metadata.request_id == request.request_id
+        assert len(metadata.sequences[0].block_ids[0]) == 4
+
+    def test_diffusion_kv_capacity_backpressures_fifo_until_blocks_are_freed(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=3, max_num_seqs=2)
+        first_request = _make_request("first")
+        second_request = _make_request("second")
+        _attach_diffusion_kv(first_request)
+        _attach_diffusion_kv(second_request)
+        self.scheduler.add_request(first_request)
+        self.scheduler.add_request(second_request)
+
+        first_output = self.scheduler.schedule()
+
+        assert _new_ids(first_output) == ["first"]
+        assert first_output.num_waiting_reqs == 1
+
+        self.scheduler.update_from_output(first_output, _make_request_output("first"))
+        second_output = self.scheduler.schedule()
+
+        assert _new_ids(second_output) == ["second"]
+        assert second_output.num_waiting_reqs == 0
+        # The same output carries both the release and the newly admitted block
+        # table; the future Worker data plane must process finished ids first.
+        assert second_output.finished_req_ids == {"first"}
+        assert second_output.scheduled_new_reqs[0].diffusion_kv_metadata is not None
+
+    def test_impossible_diffusion_kv_capacity_finishes_only_that_request(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=2, max_num_seqs=2)
+        impossible = _make_request("impossible")
+        schedulable = _make_request("schedulable")
+        _attach_diffusion_kv(impossible, seq_len=8)
+        schedulable.diffusion_kv_requests = (
+            DiffusionKVRequest(
+                "schedulable/diffusion-kv/0",
+                sequence_id=0,
+                prefix_len=0,
+                target_len=4,
+                seq_len=4,
+            ),
+        )
+        self.scheduler.add_request(impossible)
+        self.scheduler.add_request(schedulable)
+
+        sched_output = self.scheduler.schedule()
+
+        failed_state = self.scheduler.get_request_state("impossible")
+        assert failed_state is not None
+        assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert failed_state.error is not None
+        assert "cannot fit even when the block pool is empty" in failed_state.error
+        assert sched_output.finished_req_ids == {"impossible"}
+        assert _new_ids(sched_output) == ["schedulable"]
+
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            _make_request_output("schedulable"),
+        )
+
+        assert finished == {"impossible", "schedulable"}
+
+    def test_impossible_diffusion_kv_capacity_does_not_block_waiters_under_load(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=4, max_num_seqs=2)
+        running = _make_request("running")
+        _attach_diffusion_kv(running)
+        self.scheduler.add_request(running)
+        assert _new_ids(self.scheduler.schedule()) == ["running"]
+
+        impossible = _make_request("impossible")
+        impossible.diffusion_kv_requests = tuple(
+            DiffusionKVRequest(
+                f"impossible/diffusion-kv/{sequence_id}",
+                sequence_id=sequence_id,
+                prefix_len=0,
+                target_len=4,
+                seq_len=4,
+            )
+            for sequence_id in range(4)
+        )
+        schedulable = _make_request("schedulable")
+        schedulable.diffusion_kv_requests = (
+            DiffusionKVRequest(
+                "schedulable/diffusion-kv/0",
+                sequence_id=0,
+                prefix_len=0,
+                target_len=4,
+                seq_len=4,
+            ),
+        )
+        self.scheduler.add_request(impossible)
+        self.scheduler.add_request(schedulable)
+
+        sched_output = self.scheduler.schedule()
+
+        failed_state = self.scheduler.get_request_state("impossible")
+        assert failed_state is not None
+        assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert failed_state.error is not None
+        assert "required_blocks=4, available_blocks=3" in failed_state.error
+        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:
+        _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
+
+        def raise_native_error(*args, **kwargs):
+            raise ValueError("native allocation bug")
+
+        monkeypatch.setattr(
+            manager,
+            "reserve_request",
+            raise_native_error,
+        )
+
+        sched_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 == []
+
+    def test_diffusion_kv_preemption_retains_allocation(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=3)
+        request = _make_request("preempted")
+        _attach_diffusion_kv(request)
+        self.scheduler.add_request(request)
+        self.scheduler.schedule()
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+        metadata = manager.get_metadata("preempted")
+
+        assert self.scheduler.preempt_request("preempted") is True
+        resumed_output = self.scheduler.schedule()
+
+        assert resumed_output.scheduled_new_reqs == []
+        assert _cached_ids(resumed_output) == ["preempted"]
+        assert manager.get_metadata("preempted") == metadata
+
+    @pytest.mark.parametrize(
+        "status",
+        [
+            DiffusionRequestStatus.FINISHED_COMPLETED,
+            DiffusionRequestStatus.FINISHED_ABORTED,
+            DiffusionRequestStatus.FINISHED_ERROR,
+        ],
+    )
+    def test_diffusion_kv_terminal_status_releases_allocation(self, status: DiffusionRequestStatus) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=3)
+        request = _make_request("terminal")
+        _attach_diffusion_kv(request)
+        self.scheduler.add_request(request)
+        self.scheduler.schedule()
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+        assert manager.has_request("terminal") is True
+
+        self.scheduler.finish_requests("terminal", status)
+
+        assert manager.has_request("terminal") is False
 
     def test_paged_scheduler_rejects_missing_kv_request(self) -> None:
-        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        _initialize_paged_scheduler(self.scheduler)
         request = _make_request("diffusion-kv")
 
         with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"):
@@ -443,7 +656,7 @@ class TestRequestScheduler:
         ],
     )
     def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None:
-        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        _initialize_paged_scheduler(self.scheduler)
         request = _make_request("diffusion-kv")
         owner = request if owner_name == "request" else request.sampling_params
         setattr(owner, field_name, object())
@@ -921,6 +1134,45 @@ class TestDiffusionEngine:
         assert engine.scheduler.max_num_running_reqs == 1
         fake_executor_cls.assert_called_once_with(od_config)
 
+    def test_initializes_paged_scheduler_from_native_control_plane(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        mocker: MockerFixture,
+    ) -> None:
+        od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
+        fake_executor_cls = mocker.Mock(return_value=mocker.Mock())
+        scheduler = mocker.Mock()
+        kv_cache_config = object()
+        kv_vllm_config = object()
+
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+            lambda *args, **kwargs: fake_executor_cls,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+            lambda *args, **kwargs: (kv_cache_config, 16, 16, kv_vllm_config),
+        )
+
+        engine = DiffusionEngine(od_config, scheduler=scheduler)
+
+        assert engine.scheduler is scheduler
+        scheduler.initialize.assert_called_once_with(
+            od_config,
+            kv_cache_config=kv_cache_config,
+            scheduler_block_size=16,
+            hash_block_size=16,
+            kv_vllm_config=kv_vllm_config,
+        )
+
     def test_scheduler_alias_keeps_default_request_scheduler(self) -> None:
         scheduler = Scheduler()
         scheduler.initialize(SimpleNamespace())
@@ -987,6 +1239,21 @@ class TestDiffusionEngine:
         assert output.aborted is True
         assert output.abort_message == "Request req-finalize aborted."
 
+    def test_finalize_finished_request_returns_scheduler_admission_error(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = StepScheduler()
+        engine.scheduler.initialize(SimpleNamespace())
+
+        req_id = engine.scheduler.add_request(_make_request("req-error"))
+        engine.scheduler._finish_requests(
+            {req_id: DiffusionRequestStatus.FINISHED_ERROR},
+            {req_id: "KV request cannot fit"},
+        )
+
+        output = engine._finalize_finished_request(req_id)
+
+        assert output.error == "KV request cannot fit"
+
     @pytest.mark.asyncio
     async def test_streaming_runner_output_notifies_each_chunk(self) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)

tests/diffusion/test_diffusion_worker.py

@@ -13,11 +13,7 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 
-from vllm_omni.diffusion.worker.diffusion_worker import (
-    DiffusionWorker,
-    _create_diffusion_worker_vllm_config,
-    _make_diffusion_vllm_model_config,
-)
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 
@@ -58,34 +54,6 @@ def mock_gpu_worker(mocker: MockerFixture, mock_od_config):
     return worker
 
 
-def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods():
-    from types import SimpleNamespace
-
-    from vllm_omni.quantization import build_quant_config
-
-    od_config = SimpleNamespace(
-        model="dummy",
-        dtype=torch.bfloat16,
-        quantization_config=build_quant_config(
-            {
-                "quant_method": "modelopt",
-                "quant_algo": "FP8",
-                "ignore": [],
-            }
-        ),
-        tf_model_config=SimpleNamespace(),
-        enforce_eager=True,
-        is_moe=False,
-    )
-
-    model_config = _make_diffusion_vllm_model_config(od_config)
-
-    assert model_config.dtype is torch.bfloat16
-    assert model_config.quantization == "modelopt"
-    assert model_config.quantization_config is od_config.quantization_config
-    assert model_config.is_quantized()
-
-
 class TestDiffusionWorkerSleep:
     """Test DiffusionWorker.sleep method."""
 
@@ -329,49 +297,3 @@ class TestDiffusionWorkerWakeUp:
         mock_buffer2.data.copy_.assert_not_called()
 
         assert result is True
-
-
-class TestWorkerVllmConfigAdditionalConfig:
-    """Test worker-side VllmConfig construction with additional_config."""
-
-    def test_create_diffusion_worker_vllm_config_passes_additional_config(self, mocker: MockerFixture):
-        mock_vllm_config = mocker.Mock()
-        mock_vllm_cls = mocker.patch(
-            "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
-            return_value=mock_vllm_config,
-        )
-        od_config = mocker.Mock(additional_config={"torchair_graph_config": {"enabled": True}})
-
-        result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
-
-        assert result is mock_vllm_config
-        assert mock_vllm_cls.call_args.kwargs["additional_config"] == od_config.additional_config
-
-    def test_create_diffusion_worker_vllm_config_falls_back_when_constructor_rejects_additional_config(
-        self, mocker: MockerFixture
-    ):
-        mock_vllm_config = mocker.Mock()
-        mock_vllm_cls = mocker.patch(
-            "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
-            side_effect=[
-                TypeError("VllmConfig.__init__() got an unexpected keyword argument 'additional_config'"),
-                mock_vllm_config,
-            ],
-        )
-        od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
-
-        result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
-
-        assert result is mock_vllm_config
-        assert mock_vllm_cls.call_count == 2
-        assert getattr(mock_vllm_config, "additional_config") == od_config.additional_config
-
-    def test_create_diffusion_worker_vllm_config_reraises_other_type_errors(self, mocker: MockerFixture):
-        mocker.patch(
-            "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
-            side_effect=TypeError("additional_config values must be hashable"),
-        )
-        od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
-
-        with pytest.raises(TypeError, match="hashable"):
-            _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -413,6 +413,26 @@ class TestRequestModeDispatch:
         assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
         executor.collective_rpc.assert_called_once()
 
+    def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+        scheduler_output = _make_sched_output("A", "B")
+        for new_req in scheduler_output.scheduled_new_reqs:
+            new_req.diffusion_kv_metadata = SimpleNamespace(request_id=new_req.request_id)
+
+        result = executor.execute_request(scheduler_output)
+
+        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+        forwarded_envelopes = executor.collective_rpc.call_args.kwargs["args"][0]
+        assert forwarded_envelopes is scheduler_output.scheduled_new_reqs
+        assert [envelope.diffusion_kv_metadata.request_id for envelope in forwarded_envelopes] == ["A", "B"]
+
     @pytest.mark.parametrize(
         ("field", "value"),
         [

vllm_omni/config/omni_config.py

@@ -814,6 +814,10 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset(
         "omni_kv_config",
         "cfg_kv_collect_func",
         "max_num_seqs",
+        "kv_cache_memory_bytes",
+        "gpu_memory_utilization",
+        "max_num_batched_tokens",
+        "max_model_len",
         "enable_sleep_mode",
         "enforce_eager",
         "enable_multithread_weight_load",

vllm_omni/diffusion/attention/backends/abstract.py

@@ -81,6 +81,19 @@ class AttentionBackend(ABC):
         supported_head_sizes = cls.get_supported_head_sizes()
         return (not supported_head_sizes) or head_size in supported_head_sizes
 
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        """Whether this backend reads K/V pages by the runtime block stride.
+
+        Returning ``True`` means the physical cache layout has ``num_blocks``
+        as its outer stride, so native vLLM may safely use page-size padding
+        when it unifies cache layouts across layers. Dense diffusion backends
+        conservatively keep the default ``False``; a paged backend should
+        override this only when its kernel actually follows that layout.
+        """
+
+        return False
+
 
 @dataclass(frozen=True, slots=True)
 class QueryRange:

vllm_omni/diffusion/attention/layer.py

@@ -11,8 +11,10 @@ from dataclasses import replace
 
 import torch
 import torch.nn as nn
+from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.logger import init_logger
 from vllm.model_executor.models.utils import extract_layer_index
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
@@ -61,12 +63,23 @@ class Attention(nn.Module):
         # perf for this layer (e.g. Wan2.2 cross-attn has short sequences and
         # block-FP8 quant offers no win). Default False = follow global config.
         disable_kv_quant: bool = False,
+        # Opt-in marker for Scheduler-managed paged KV. Unmarked diffusion
+        # attention remains dense and contributes no native KVCacheSpec.
+        paged_kv_cache_role: str | None = None,
+        paged_kv_cache_dtype: torch.dtype | None = None,
     ):
         super().__init__()
 
         self.role = role
         self.role_category = role_category
         self.qkv_layout = qkv_layout
+        if paged_kv_cache_role == "":
+            raise ValueError("paged_kv_cache_role must be non-empty when provided")
+        self.paged_kv_cache_role = paged_kv_cache_role
+        self.paged_kv_cache_dtype = paged_kv_cache_dtype
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.head_size = head_size
 
         # Resolve backend via role-aware config.
         # The global diffusion config is set during model init via
@@ -171,6 +184,25 @@ class Attention(nn.Module):
         self._disable_kv_quant: bool = disable_kv_quant
         self._init_kv_cache_quantization(config)
 
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
+        """Return native rank-local geometry for an opted-in paged cache."""
+
+        if self.paged_kv_cache_role is None:
+            return None
+        dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype
+        # Keep backend layout discovery under the same config context used by
+        # upstream vLLM's attention-spec collector.
+        with set_current_vllm_config(vllm_config):
+            indexes_kv_by_block_stride = self.attn_backend.indexes_kv_by_block_stride()
+        return FullAttentionSpec(
+            block_size=vllm_config.cache_config.block_size,
+            num_kv_heads=self.num_kv_heads,
+            head_size=self.head_size,
+            dtype=dtype,
+            indexes_kv_by_block_stride=indexes_kv_by_block_stride,
+            non_causal=not self.causal,
+        )
+
     def _get_active_parallel_strategy(self):
         """Get the parallel strategy based on current SP active state.
 

vllm_omni/diffusion/data.py

@@ -851,6 +851,14 @@ class OmniDiffusionConfig:
     # Maximum number of sequences to generate in a batch
     max_num_seqs: int = 1
 
+    # Native vLLM KV-cache sizing inputs used only by paged_scheduler mode.
+    # These mirror the structured stage cache/scheduler config until all
+    # diffusion runtime construction consumes VllmOmniConfig directly.
+    kv_cache_memory_bytes: int | None = None
+    gpu_memory_utilization: float = 0.9
+    max_num_batched_tokens: int | None = None
+    max_model_len: int | None = None
+
     # Request-mode batch admission: wait briefly for compatible requests to
     # accumulate in the scheduler waiting queue before the first schedule() of
     # a wave.  Improves fused forward batch sizes under bursty HTTP ingress.
@@ -936,6 +944,14 @@ class OmniDiffusionConfig:
         if not isinstance(self.diffusion_compile_dynamic, bool):
             raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
         self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
+        if self.kv_cache_memory_bytes is not None and self.kv_cache_memory_bytes < 0:
+            raise ValueError("kv_cache_memory_bytes must be non-negative")
+        if not 0.0 < self.gpu_memory_utilization <= 1.0:
+            raise ValueError("gpu_memory_utilization must be in (0, 1]")
+        if self.max_num_batched_tokens is not None and self.max_num_batched_tokens <= 0:
+            raise ValueError("max_num_batched_tokens must be positive")
+        if self.max_model_len is not None and self.max_model_len != -1 and self.max_model_len <= 0:
+            raise ValueError("max_model_len must be positive or -1")
 
         if self.omni_kv_config is None:
             self.omni_kv_config = {}

vllm_omni/diffusion/diffusion_engine.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures
+import copy
 import inspect
 import queue
 import threading
@@ -17,18 +18,23 @@ from typing import TYPE_CHECKING, Any
 import numpy as np
 import PIL.Image
 import torch
+from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.kv_cache_interface import KVCacheConfig
 
 from vllm_omni.diffusion.data import (
     DiffusionOutput,
     DiffusionRequestAbortedError,
     OmniDiffusionConfig,
 )
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.initialization import initialize_diffusion_kv_control_plane
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.io_support import (
     get_dummy_run_num_frames,
+    get_dummy_run_num_image_inputs,
     image_color_format,
     supports_audio_output,
     supports_multimodal_input,
@@ -184,12 +190,39 @@ class DiffusionEngine:
                 from the resolved execution mode.
         """
         self.od_config = od_config
+        # Set after the paged-KV profile request has gone through model-owned
+        # preprocessing. Real requests are admitted only within this measured
+        # activation envelope: (max execution sequences, max seq_len,
+        # max target_len).
+        self._diffusion_kv_profile_limits: tuple[int, int, int] | None = None
 
         self._init_process_hooks(od_config)
         self.execution_mode = self._resolve_execution_mode(od_config)
         self._init_executor(od_config)
         try:
-            self._init_scheduler(od_config, scheduler)
+            profile_requests = self._prepare_diffusion_kv_profile_requests()
+            kv_control_plane = initialize_diffusion_kv_control_plane(
+                self.executor,
+                od_config,
+                profile_requests=profile_requests,
+            )
+            if kv_control_plane is None:
+                self._init_scheduler(od_config, scheduler)
+            else:
+                (
+                    kv_cache_config,
+                    scheduler_block_size,
+                    hash_block_size,
+                    kv_vllm_config,
+                ) = kv_control_plane
+                self._init_scheduler(
+                    od_config,
+                    scheduler,
+                    kv_cache_config,
+                    scheduler_block_size=scheduler_block_size,
+                    hash_block_size=hash_block_size,
+                    kv_vllm_config=kv_vllm_config,
+                )
             self._init_runtime_state()
             self._init_execute_fn()
             self._log_execution_mode(od_config)
@@ -242,6 +275,11 @@ class DiffusionEngine:
         self,
         od_config: OmniDiffusionConfig,
         scheduler: BaseScheduler | None = None,
+        kv_cache_config: KVCacheConfig | None = None,
+        *,
+        scheduler_block_size: int | None = None,
+        hash_block_size: int | None = None,
+        kv_vllm_config: VllmConfig | None = None,
     ) -> None:
         if scheduler is not None:
             self.scheduler = scheduler
@@ -249,7 +287,16 @@ class DiffusionEngine:
             self.scheduler = StepScheduler()
         else:
             self.scheduler = RequestScheduler()
-        self.scheduler.initialize(od_config)
+        if kv_cache_config is None:
+            self.scheduler.initialize(od_config)
+        else:
+            self.scheduler.initialize(
+                od_config,
+                kv_cache_config=kv_cache_config,
+                scheduler_block_size=scheduler_block_size,
+                hash_block_size=hash_block_size,
+                kv_vllm_config=kv_vllm_config,
+            )
 
     def _init_runtime_state(self) -> None:
         # DP multi-concurrency: allow batching dp_size requests so each
@@ -703,9 +750,36 @@ class DiffusionEngine:
         """Run model-owned preprocessing once, before entering Engine locks."""
 
         pre_process_func = getattr(self, "pre_process_func", None)
-        if pre_process_func is None:
-            return request
-        return pre_process_func(request)
+        if pre_process_func is not None:
+            request = pre_process_func(request)
+        self._validate_diffusion_kv_profile_limits(request)
+        return request
+
+    def _validate_diffusion_kv_profile_limits(self, request: OmniDiffusionRequest) -> None:
+        """Keep admitted paged-KV requests within the profiled activation shape."""
+
+        profile_limits = getattr(self, "_diffusion_kv_profile_limits", None)
+        if profile_limits is None:
+            return
+        kv_requests = request.diffusion_kv_requests
+        if not kv_requests:
+            return
+
+        request_limits = (
+            len(kv_requests),
+            max(kv_request.seq_len for kv_request in kv_requests),
+            max(kv_request.target_len for kv_request in kv_requests),
+        )
+        if any(actual > profiled for actual, profiled in zip(request_limits, profile_limits, strict=True)):
+            request_sequences, request_seq_len, request_target_len = request_limits
+            profile_sequences, profile_seq_len, profile_target_len = profile_limits
+            raise ValueError(
+                f"Diffusion KV request {request.request_id!r} exceeds the startup memory-profile envelope: "
+                f"sequences={request_sequences} (profiled={profile_sequences}), "
+                f"max_seq_len={request_seq_len} (profiled={profile_seq_len}), "
+                f"max_target_len={request_target_len} (profiled={profile_target_len}). "
+                "Reduce the request shape or extend the model's paged-KV profile recipe."
+            )
 
     def _add_prepared_request(self, request: OmniDiffusionRequest) -> str:
         """Admit a request whose model-owned preprocessing is complete."""
@@ -869,47 +943,111 @@ class DiffusionEngine:
             self.close()
             raise e
 
-    def _dummy_run(self):
-        """A dummy run to warm up the model."""
-        num_inference_steps = 1
-        height = 512
-        width = 512
-        prompt: OmniTextPrompt = {"prompt": "dummy run"}
+    def _make_dummy_request(
+        self,
+        *,
+        height: int,
+        width: int,
+        guidance_scale: float,
+        num_image_inputs: int = 1,
+    ) -> OmniDiffusionRequest | None:
+        """Build a one-step model request for startup profiling or warmup."""
 
+        prompt: OmniTextPrompt = {"prompt": "dummy run"}
         supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
         if supports_image_input:
-            # Provide a dummy image input if the model supports it
             color_format = image_color_format(self.od_config.model_class_name)
-            dummy_image = PIL.Image.new(color_format, (width, height))
-            prompt.setdefault("multi_modal_data", {})["image"] = dummy_image
+            images = [PIL.Image.new(color_format, (width, height)) for _ in range(num_image_inputs)]
+            prompt.setdefault("multi_modal_data", {})["image"] = images[0] if len(images) == 1 else images
 
         if supports_audio_input:
             audio_sr = 16000
-            dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
-            prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
+            prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32)
 
         num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
         if num_frames <= 0:
-            logger.info("Skipping dummy warmup run (num_frames=0)")
-            return
-        req = OmniDiffusionRequest(
+            return None
+        return OmniDiffusionRequest(
             prompt=prompt,
             request_id=DUMMY_DIFFUSION_REQUEST_ID,
             sampling_params=OmniDiffusionSamplingParams(
                 height=height,
                 width=width,
-                num_inference_steps=num_inference_steps,
+                num_inference_steps=1,
                 num_frames=num_frames,
-                # Keep warmup path minimal and robust across text encoders.
-                # Some models may fail when warmup implicitly triggers
-                # classifier-free guidance with an empty negative prompt.
-                guidance_scale=0.0,
+                guidance_scale=guidance_scale,
                 num_outputs_per_prompt=1,
-                # Disable CFG for warmup to avoid triggering CFG parallel
-                # validation when cfg_parallel_size > 1.
                 extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0},
             ),
         )
+
+    def _prepare_diffusion_kv_profile_requests(self) -> list[OmniDiffusionRequest] | None:
+        """Prepare the per-rank request batch used to profile paged-KV headroom.
+
+        The profile executes directly on each Worker before the Scheduler and
+        its KV manager exist. It uses the maximum number of requests that one
+        rank can execute together. DLO+DP is the exception because each rank
+        executes one request from the collective wave.
+
+        Hunyuan is currently the only model integrated with
+        ``paged_scheduler``; 1024x1024, 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.
+        Future paged model integrations must extend this recipe for their
+        serving limits (for example, video frame count) rather than reusing it
+        silently.
+        """
+
+        if (
+            getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            is not DiffusionKVCacheMode.PAGED_SCHEDULER
+        ):
+            return None
+        request = self._make_dummy_request(
+            height=1024,
+            width=1024,
+            guidance_scale=5.0,
+            num_image_inputs=get_dummy_run_num_image_inputs(self.od_config.model_class_name),
+        )
+        if request is None:
+            raise RuntimeError("paged_scheduler requires a runnable Diffusion KV memory profile request")
+        request = self._prepare_request_for_admission(request)
+        kv_requests = request.diffusion_kv_requests
+        if not kv_requests:
+            raise RuntimeError("paged_scheduler profile preprocessing must produce Diffusion KV requests")
+        self._diffusion_kv_profile_limits = (
+            len(kv_requests),
+            max(kv_request.seq_len for kv_request in kv_requests),
+            max(kv_request.target_len for kv_request in kv_requests),
+        )
+        dlo_dp_request_mode = self.execution_mode is DiffusionExecutionMode.REQUEST_BATCH and _uses_dlo_dp_concurrency(
+            self.od_config
+        )
+        profile_batch_size = 1 if dlo_dp_request_mode else _max_num_seqs(self.od_config)
+        profile_requests: list[OmniDiffusionRequest] = []
+        for index in range(profile_batch_size):
+            profile_request = copy.copy(request)
+            profile_request.request_id = f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-{index}"
+            profile_request.sampling_params = copy.deepcopy(request.sampling_params)
+            # Preprocessing owns request geometry and prepared model inputs,
+            # but native KV requests remain Scheduler-only mutable state. The
+            # profile bypasses Scheduler admission and must not send them to
+            # the Worker.
+            profile_request.diffusion_kv_requests = None
+            profile_requests.append(profile_request)
+        return profile_requests
+
+    def _dummy_run(self):
+        """A dummy run to warm up the model."""
+        req = self._make_dummy_request(
+            height=512,
+            width=512,
+            guidance_scale=0.0,
+        )
+        if req is None:
+            logger.info("Skipping dummy warmup run (num_frames=0)")
+            return
         logger.info("dummy run to warm up the model")
         output = self.add_req_and_wait_for_response(req)
         if output.error:
@@ -1127,4 +1265,7 @@ class DiffusionEngine:
         if runner_output is not None and runner_output.async_output_id is not None:
             return DiffusionOutput(async_output_id=runner_output.async_output_id)
 
+        if state.status == DiffusionRequestStatus.FINISHED_ERROR and state.error:
+            return DiffusionOutput(error=state.error)
+
         return DiffusionOutput(error=missing_result_error)

vllm_omni/diffusion/diffusion_kv/initialization.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from vllm.config import VllmConfig
+from vllm.v1.core.kv_cache_utils import (
+    generate_scheduler_kv_cache_config,
+    get_kv_cache_configs,
+    resolve_kv_cache_block_sizes,
+)
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
+from vllm_omni.platforms import current_omni_platform
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+
+def build_native_kv_cache_configs(
+    vllm_config: VllmConfig,
+    worker_specs: list[dict[str, KVCacheSpec]],
+    available_memory: list[int],
+) -> tuple[list[KVCacheConfig], KVCacheConfig]:
+    """Build rank-local and Scheduler-native configs using vLLM utilities."""
+
+    from vllm.v1.core import single_type_kv_cache_manager
+
+    single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config)
+    worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory)
+    return worker_configs, generate_scheduler_kv_cache_config(worker_configs)
+
+
+def initialize_diffusion_kv_control_plane(
+    executor: DiffusionExecutor,
+    od_config: OmniDiffusionConfig,
+    *,
+    profile_requests: list[OmniDiffusionRequest] | None = None,
+    device: torch.device | None = None,
+) -> tuple[KVCacheConfig, int, int, VllmConfig] | None:
+    """Run the native Worker-spec to Scheduler-config initialization chain."""
+
+    if (
+        getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+        is not DiffusionKVCacheMode.PAGED_SCHEDULER
+    ):
+        return None
+    if not profile_requests:
+        raise ValueError("paged_scheduler Diffusion KV initialization requires prepared memory profile requests")
+    if device is None:
+        device = current_omni_platform.get_torch_device(0)
+    vllm_config = create_diffusion_vllm_config(device, od_config)
+    worker_specs = executor.get_kv_cache_specs()
+    available_memory = executor.determine_available_kv_memory(profile_requests)
+    if len(worker_specs) != od_config.num_gpus or len(available_memory) != od_config.num_gpus:
+        raise ValueError(
+            "Diffusion KV initialization rank count mismatch: "
+            f"expected={od_config.num_gpus}, specs={len(worker_specs)}, memory={len(available_memory)}"
+        )
+    worker_configs, scheduler_kv_cache_config = build_native_kv_cache_configs(
+        vllm_config,
+        worker_specs,
+        available_memory,
+    )
+    scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(
+        scheduler_kv_cache_config,
+        vllm_config,
+    )
+    executor.set_kv_cache_configs(worker_configs)
+    return (
+        scheduler_kv_cache_config,
+        scheduler_block_size,
+        hash_block_size,
+        vllm_config,
+    )

vllm_omni/diffusion/diffusion_kv/manager.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+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.diffusion.diffusion_kv.metadata import (
+    DiffusionKVMetadata,
+    DiffusionKVSequenceMetadata,
+)
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+
+
+class DiffusionKVAdmissionError(ValueError):
+    """A request-specific KV admission failure safe to return to the caller."""
+
+
+class DiffusionKVCacheManager:
+    """Atomic public-request facade over native vLLM KVCacheManager."""
+
+    def __init__(
+        self,
+        kv_cache_config: KVCacheConfig,
+        *,
+        max_model_len: int,
+        scheduler_block_size: int,
+        hash_block_size: int,
+        max_in_flight_tokens: int | None = None,
+    ) -> None:
+        if max_model_len <= 0:
+            raise ValueError(f"max_model_len must be positive, got {max_model_len}")
+        if kv_cache_config.num_blocks <= 0 or not kv_cache_config.kv_cache_groups:
+            raise ValueError("Diffusion KVCacheConfig must contain a positive block pool and at least one group")
+        if scheduler_block_size <= 0 or hash_block_size <= 0:
+            raise ValueError("scheduler_block_size and hash_block_size must be positive")
+        if max_in_flight_tokens is not None and max_in_flight_tokens <= 0:
+            raise ValueError(f"max_in_flight_tokens must be positive when provided, got {max_in_flight_tokens}")
+        self.native_manager = KVCacheManager(
+            kv_cache_config=kv_cache_config,
+            max_model_len=max_model_len,
+            scheduler_block_size=scheduler_block_size,
+            hash_block_size=hash_block_size,
+            max_in_flight_tokens=max_in_flight_tokens,
+            enable_caching=False,
+        )
+        self.max_model_len = max_model_len
+        # Native vLLM may reserve a null block, so an idle BlockPool does not
+        # necessarily report ``kv_cache_config.num_blocks`` free blocks.
+        self._empty_pool_num_free_blocks = self.native_manager.block_pool.get_num_free_blocks()
+        self._requests: dict[str, tuple[DiffusionKVRequest, ...]] = {}
+        self._metadata: dict[str, DiffusionKVMetadata] = {}
+        self._internal_request_ids: set[str] = set()
+        self._next_allocation_generation = 1
+
+    def _get_empty_pool_required_blocks(self, requests: Sequence[DiffusionKVRequest]) -> int:
+        """Return the native full-sequence reservation for one public request.
+
+        vLLM performs this calculation inside ``allocate_slots`` when
+        ``full_sequence_must_fit`` is enabled. Diffusion CFG adds one semantic
+        requirement that native requests do not have: every execution sequence
+        must fit atomically. Sum the native per-sequence result so a request
+        that can never fit is rejected independently of current pool usage.
+        """
+
+        coordinator = self.native_manager.coordinator
+        empty_blocks = self.native_manager.empty_kv_cache_blocks.blocks
+        return sum(
+            coordinator.get_num_blocks_to_allocate(
+                request_id=request.request_id,
+                num_tokens=request.num_tokens,
+                new_computed_blocks=empty_blocks,
+                num_encoder_tokens=0,
+                total_computed_tokens=0,
+                num_local_computed_tokens=0,
+                num_tokens_main_model=request.num_tokens,
+                apply_admission_cap=True,
+            )
+            for request in requests
+        )
+
+    def has_request(self, public_request_id: str) -> bool:
+        return public_request_id in self._requests
+
+    def reserve_request(
+        self,
+        public_request_id: str,
+        kv_requests: Sequence[DiffusionKVRequest],
+    ) -> DiffusionKVMetadata | None:
+        """Reserve all CFG sequences or roll back the complete public request."""
+
+        if not public_request_id:
+            raise ValueError("public_request_id must be non-empty")
+        if public_request_id in self._requests:
+            raise ValueError(f"Diffusion KV request {public_request_id!r} is already allocated")
+        requests = tuple(kv_requests)
+        if not requests:
+            raise ValueError("Diffusion KV allocation requires at least one sequence")
+
+        sequence_ids = [request.sequence_id for request in requests]
+        internal_ids = [request.request_id for request in requests]
+        expected_sequence_ids = list(range(len(requests)))
+        if sequence_ids != expected_sequence_ids:
+            raise ValueError(
+                "Diffusion KV requests must be ordered by contiguous sequence_id values: "
+                f"expected={expected_sequence_ids}, got={sequence_ids}"
+            )
+        if len(internal_ids) != len(set(internal_ids)):
+            raise ValueError(f"Diffusion KV internal request IDs must be unique, got {internal_ids}")
+        conflicts = self._internal_request_ids.intersection(internal_ids)
+        if conflicts:
+            raise ValueError(f"Diffusion KV internal request IDs are already allocated: {sorted(conflicts)}")
+        for request in requests:
+            if request.kv_contexts:
+                raise DiffusionKVAdmissionError("independent DiffusionKVContext allocation is not implemented yet")
+            if request.seq_len > self.max_model_len:
+                raise DiffusionKVAdmissionError(
+                    f"Diffusion KV sequence {request.request_id!r} exceeds max_model_len: "
+                    f"seq_len={request.seq_len}, max_model_len={self.max_model_len}"
+                )
+
+        required_blocks = self._get_empty_pool_required_blocks(requests)
+        if required_blocks > self._empty_pool_num_free_blocks:
+            raise DiffusionKVAdmissionError(
+                f"Diffusion KV request {public_request_id!r} cannot fit even when the block pool is empty: "
+                f"required_blocks={required_blocks}, available_blocks={self._empty_pool_num_free_blocks}; "
+                "increase KV cache capacity or reduce the request sequence count/length"
+            )
+
+        allocated: list[DiffusionKVRequest] = []
+        sequence_metadata: list[DiffusionKVSequenceMetadata] = []
+        try:
+            for request in requests:
+                blocks = self.native_manager.allocate_slots(
+                    request,
+                    num_new_tokens=request.seq_len,
+                    delay_cache_blocks=True,
+                    full_sequence_must_fit=True,
+                )
+                if blocks is None:
+                    self._rollback(allocated)
+                    allocated.clear()
+                    return None
+                allocated.append(request)
+                sequence_metadata.append(
+                    DiffusionKVSequenceMetadata(
+                        sequence_id=request.sequence_id,
+                        prefix_len=request.prefix_len,
+                        target_len=request.target_len,
+                        seq_len=request.seq_len,
+                        block_ids=blocks.get_block_ids(),
+                    )
+                )
+        except Exception:
+            self._rollback(allocated)
+            raise
+
+        metadata = DiffusionKVMetadata(
+            request_id=public_request_id,
+            allocation_generation=self._next_allocation_generation,
+            sequences=tuple(sequence_metadata),
+        )
+        self._next_allocation_generation += 1
+        self._requests[public_request_id] = requests
+        self._metadata[public_request_id] = metadata
+        self._internal_request_ids.update(internal_ids)
+        return metadata
+
+    def get_metadata(self, public_request_id: str) -> DiffusionKVMetadata:
+        try:
+            return self._metadata[public_request_id]
+        except KeyError as exc:
+            raise KeyError(f"Diffusion KV request {public_request_id!r} is not allocated") from exc
+
+    def free_request(self, public_request_id: str) -> None:
+        requests = self._requests.pop(public_request_id, ())
+        self._metadata.pop(public_request_id, None)
+        for request in reversed(requests):
+            self.native_manager.free(request)
+            self._internal_request_ids.discard(request.request_id)
+
+    def close(self) -> None:
+        for public_request_id in tuple(self._requests):
+            self.free_request(public_request_id)
+
+    def _rollback(self, requests: Sequence[DiffusionKVRequest]) -> None:
+        for request in reversed(requests):
+            self.native_manager.free(request)

vllm_omni/diffusion/diffusion_kv/request.py

@@ -105,7 +105,7 @@ class DiffusionKVRequest:
         self.status = RequestStatus.WAITING
         self.num_preemptions = 0
         self.num_in_flight_tokens = 0
-        # vLLM 0.26 uses this optional boundary when publishing cached blocks;
+        # vLLM 0.26+ uses this optional boundary when publishing cached blocks;
         # zero means that no sparse-retention boundary is active.
         self.shared_prefix_boundary = 0
 

vllm_omni/diffusion/executor/abstract.py

@@ -5,10 +5,12 @@ from collections.abc import Callable
 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_omni.diffusion.data import OmniDiffusionConfig
 
 if TYPE_CHECKING:
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
     from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
     from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
 
@@ -115,6 +117,39 @@ class DiffusionExecutor(ABC):
         """
         return None
 
+    def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
+        """Collect rank-local native specs after every Worker loads its model."""
+
+        result = self.collective_rpc(
+            "get_kv_cache_specs",
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        if not isinstance(result, list):
+            raise TypeError(f"get_kv_cache_specs returned {type(result).__name__}, expected list")
+        return result
+
+    def determine_available_kv_memory(self, profile_requests: list[OmniDiffusionRequest]) -> list[int]:
+        """Profile and collect the KV memory budget on every Worker rank."""
+
+        result = self.collective_rpc(
+            "determine_available_kv_memory",
+            args=(profile_requests,),
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        if not isinstance(result, list) or not all(isinstance(value, int) for value in result):
+            raise TypeError("determine_available_kv_memory must return list[int]")
+        return result
+
+    def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig]) -> None:
+        """Send each native rank-local config back to all Workers."""
+
+        # The default control-plane RPC mode executes on every rank and has
+        # rank 0 return the gathered rank statuses, so failures on nonzero
+        # ranks are not silently dropped.
+        self.collective_rpc("set_kv_cache_configs", args=(kv_cache_configs,))
+
     @abstractmethod
     def shutdown(self) -> None:
         """Shutdown the executor and release resources."""

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -465,16 +465,14 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         for new_req in new_reqs:
             validate_new_request_data_identity(new_req)
 
-        has_diffusion_kv_metadata = any(new_req.diffusion_kv_metadata is not None for new_req in new_reqs)
-
         # DP multi-concurrency: when DLO+AllGather is active and multiple
-        # requests are scheduled, send ALL requests in one broadcast RPC.
-        # Each rank picks req[rank % len(reqs)] and computes independently.
+        # requests are scheduled, send every complete NewRequestData envelope
+        # in one broadcast RPC. Each rank picks one envelope, keeping its
+        # request and Diffusion KV metadata inseparable.
         # All ranks reply (unique_reply_rank=None) so we collect dp_size
         # responses and match by dp_rank.
         if (
             len(new_reqs) > 1
-            and not has_diffusion_kv_metadata
             and getattr(self.od_config, "enable_distributed_layerwise_offload", False)
             and getattr(self.od_config, "dlo_use_allgather", True)
         ):
@@ -507,12 +505,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                     f"empty prompt request IDs: {empty_prompt_ids}."
                 )
 
-            reqs_list = [nr.req for nr in new_reqs]
             try:
                 results = self.collective_rpc(
                     "execute_model",
                     timeout=_DLO_DP_WAVE_TIMEOUT_S,
-                    args=(reqs_list, self.od_config, scheduler_output.kv_prefetch_job),
+                    args=(new_reqs, self.od_config, scheduler_output.kv_prefetch_job),
                     unique_reply_rank=None,
                     exec_all_ranks=True,
                 )

vllm_omni/diffusion/io_support.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 import inspect
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 from vllm_omni.diffusion.registry import DiffusionModelRegistry
 
 
@@ -47,3 +48,9 @@ def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool)
     if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"):
         return int(getattr(model_cls, "dummy_run_num_frames"))
     return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1
+
+
+def get_dummy_run_num_image_inputs(model_class_name: str) -> int:
+    """Return the maximum advertised image-input count for profiling."""
+
+    return get_diffusion_model_metadata(model_class_name).max_multimodal_image_inputs or 1

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -913,7 +913,7 @@ class HunYuanRotary2DEmbedder:
         return q, k
 
 
-class ImageKVCacheManager:
+class ImageKVCacheManager(nn.Module):
     """
     Manages specialized caching and updating of KV-Cache for image tokens in multimodal models.
     """
@@ -932,6 +932,7 @@ class ImageKVCacheManager:
             image_token_len: Number of tokens per image (including special placeholders),
             default 4097 (timestamp + 4096 image tokens).
         """
+        super().__init__()
         # attention related
         self.num_heads = num_heads
         self.num_kv_heads = num_kv_heads
@@ -953,6 +954,10 @@ class ImageKVCacheManager:
             softmax_scale=self.scaling,
             num_kv_heads=self.num_kv_heads,
             prefix=f"{prefix}.attn" if prefix else "",
+            paged_kv_cache_role="primary",
+            # Hunyuan image RoPE produces bfloat16 K and denoising runs under
+            # bfloat16 autocast, independent of the weight-loading dtype.
+            paged_kv_cache_dtype=torch.bfloat16,
         )
 
     @staticmethod
@@ -1095,7 +1100,7 @@ class ImageKVCacheManager:
         )
         return key, value
 
-    def __call__(
+    def forward(
         self,
         query: torch.Tensor,
         key: torch.Tensor,

vllm_omni/diffusion/request.py

@@ -14,6 +14,7 @@ if TYPE_CHECKING:
     from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 
 DUMMY_DIFFUSION_REQUEST_ID = "dummy_req_id"
+_DUMMY_DIFFUSION_REQUEST_ID_PREFIX = f"{DUMMY_DIFFUSION_REQUEST_ID}/"
 
 
 @dataclass
@@ -90,4 +91,6 @@ class OmniDiffusionRequest:
 
     @classmethod
     def is_dummy_run_request_id(cls, request_id: str | None) -> bool:
-        return request_id == DUMMY_DIFFUSION_REQUEST_ID
+        return request_id == DUMMY_DIFFUSION_REQUEST_ID or (
+            isinstance(request_id, str) and request_id.startswith(_DUMMY_DIFFUSION_REQUEST_ID_PREFIX)
+        )

vllm_omni/diffusion/sched/base_scheduler.py

@@ -7,10 +7,14 @@ from abc import ABC, abstractmethod
 from collections import deque
 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_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.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
     CachedRequestData,
@@ -49,9 +53,17 @@ class BaseScheduler(ABC):
         self._finished_req_ids: set[str] = set()
         self.max_num_running_reqs: int = 1
         self._prefetch_enabled: bool = False
-        self._diffusion_kv_enabled = False
+        self._diffusion_kv_manager: DiffusionKVCacheManager | None = None
 
-    def initialize(self, od_config: OmniDiffusionConfig) -> None:
+    def initialize(
+        self,
+        od_config: OmniDiffusionConfig,
+        *,
+        kv_cache_config: KVCacheConfig | None = None,
+        scheduler_block_size: int | None = None,
+        hash_block_size: int | None = None,
+        kv_vllm_config: VllmConfig | None = None,
+    ) -> None:
         self.od_config = od_config
         self._request_states.clear()
         self._step_id = 0
@@ -66,10 +78,36 @@ class BaseScheduler(ABC):
             self.max_num_running_reqs = 1
         omni_kv = getattr(od_config, "omni_kv_config", None) or {}
         self._prefetch_enabled = bool(omni_kv.get("enable_kv_async_prefetch", False))
-        self._diffusion_kv_enabled = (
+        diffusion_kv_enabled = (
             getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
             is DiffusionKVCacheMode.PAGED_SCHEDULER
         )
+        if diffusion_kv_enabled:
+            if kv_cache_config is None:
+                raise ValueError("paged_scheduler Diffusion KV requires a native Scheduler KVCacheConfig")
+            if scheduler_block_size is None or hash_block_size is None:
+                raise ValueError("paged_scheduler Diffusion KV requires native scheduler/hash block sizes")
+            if kv_vllm_config is None:
+                raise ValueError("paged_scheduler Diffusion KV requires the native VllmConfig used for cache sizing")
+            self._diffusion_kv_manager = DiffusionKVCacheManager(
+                kv_cache_config,
+                max_model_len=kv_vllm_config.model_config.max_model_len,
+                scheduler_block_size=scheduler_block_size,
+                hash_block_size=hash_block_size,
+                max_in_flight_tokens=kv_vllm_config.max_in_flight_tokens,
+            )
+        else:
+            if any(
+                value is not None
+                for value in (
+                    kv_cache_config,
+                    scheduler_block_size,
+                    hash_block_size,
+                    kv_vllm_config,
+                )
+            ):
+                raise ValueError("dense_legacy Scheduler received unexpected Diffusion KV cache initialization state")
+            self._diffusion_kv_manager = None
         self._reset_scheduler_state()
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
@@ -104,6 +142,35 @@ class BaseScheduler(ABC):
             if not self._can_schedule_waiting(state):
                 break
 
+            diffusion_kv_metadata: DiffusionKVMetadata | None = None
+            if self._diffusion_kv_manager is not None:
+                if self._diffusion_kv_manager.has_request(request_id):
+                    diffusion_kv_metadata = self._diffusion_kv_manager.get_metadata(request_id)
+                else:
+                    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,
+                            )
+                        self._finish_requests(
+                            {request_id: DiffusionRequestStatus.FINISHED_ERROR},
+                            {request_id: str(exc)},
+                        )
+                        continue
+                    if allocation is None:
+                        break
+                    diffusion_kv_metadata = allocation
+
             self._waiting.popleft()
             was_new_request = state.status == DiffusionRequestStatus.WAITING
             if not self._running:
@@ -111,7 +178,12 @@ class BaseScheduler(ABC):
             state.status = DiffusionRequestStatus.RUNNING
             self._running.append(request_id)
             if was_new_request:
-                scheduled_new_reqs.append(NewRequestData.from_state(state))
+                scheduled_new_reqs.append(
+                    NewRequestData.from_state(
+                        state,
+                        diffusion_kv_metadata=diffusion_kv_metadata,
+                    )
+                )
             else:
                 scheduled_cached_request_ids.append(request_id)
 
@@ -184,6 +256,8 @@ class BaseScheduler(ABC):
         return self._request_states.get(request_id)
 
     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)
         return self._request_states.pop(request_id, None)
 
@@ -206,6 +280,9 @@ 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
         self._request_states.clear()
         self._waiting.clear()
         self._running.clear()
@@ -245,6 +322,8 @@ 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]
             state.status = status
@@ -266,7 +345,13 @@ class BaseScheduler(ABC):
         # update_from_output() processes the runner output. It is already
         # marked finished at that point, but we still need to surface its id
         # in this update so the engine can observe the terminal state.
+        # Also surface admission failures recorded while schedule() built this
+        # output. Older finished ids retained only for Worker cleanup have
+        # already been popped by the Engine and are deliberately ignored.
         finished_req_ids = {
+            request_id for request_id in sched_output.finished_req_ids if request_id in self._request_states
+        }
+        finished_req_ids |= {
             request_id for request_id in sched_output.scheduled_request_ids if request_id in self._finished_req_ids
         }
         finished_req_ids |= self._finish_requests(statuses, errors)
@@ -280,7 +365,7 @@ class BaseScheduler(ABC):
 
     def _make_request_state(self, request_id: str, request: OmniDiffusionRequest) -> SchedulerRequestState:
         kv_requests = request.diffusion_kv_requests or ()
-        if self._diffusion_kv_enabled:
+        if self._diffusion_kv_manager is not None:
             self._reject_legacy_dense_kv(request)
             if not kv_requests:
                 raise ValueError("paged_scheduler request preprocessing did not produce DiffusionKVRequest state")

vllm_omni/diffusion/vllm_config.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Build native vLLM configuration views for the Omni diffusion runtime."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.config import CompilationConfig, DeviceConfig, VllmConfig
+from vllm.transformers_utils.config import get_hf_text_config
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+
+
+def resolve_diffusion_max_model_len(od_config: OmniDiffusionConfig) -> int:
+    """Resolve the native KV admission bound for a diffusion stage."""
+
+    configured = getattr(od_config, "max_model_len", None)
+    if configured is not None:
+        configured = int(configured)
+        if configured == -1:
+            configured = None
+        elif configured <= 0:
+            raise ValueError(f"max_model_len must be positive for Scheduler-managed Diffusion KV, got {configured}")
+        else:
+            return configured
+
+    # ``max_model_len`` keeps the native vLLM meaning: the per-sequence KV
+    # admission ceiling consumed by ``KVCacheManager``. HunyuanImage3 exposes
+    # that ceiling directly as ``max_position_embeddings`` because its text,
+    # reference-image, and generated-image tokens share one Transformer
+    # sequence. Other diffusion models commonly express their serving limits
+    # through resolution, duration, or frame-count constraints instead. When
+    # paged KV support expands to those models, the model integration must
+    # derive ``max_model_len`` from its serving limits, while model-owned
+    # preprocessing validates each request and derives its matching
+    # ``seq_len``. Text-encoder length fields must not be mistaken for the DiT
+    # KV limit.
+    hf_config = getattr(od_config, "tf_model_config", None)
+    for field_name in ("max_position_embeddings", "max_sequence_length", "seq_length"):
+        value = getattr(hf_config, field_name, None)
+        if value is None and isinstance(hf_config, dict):
+            value = hf_config.get(field_name)
+        if value is not None and int(value) > 0:
+            return int(value)
+
+    mode = getattr(getattr(od_config, "diffusion_kv_mode", None), "value", None)
+    if mode == "paged_scheduler":
+        raise ValueError(
+            "paged_scheduler Diffusion KV requires max_model_len or a model config "
+            "with max_position_embeddings/max_sequence_length/seq_length"
+        )
+    # Dense execution does not consume the native KV admission bound, but the
+    # shared Worker VllmConfig still requires a positive placeholder.
+    return 1
+
+
+@dataclass
+class _DiffusionVllmModelConfig:
+    """Private adapter for the native ``VllmConfig.model_config`` surface."""
+
+    model: str | None
+    dtype: torch.dtype
+    max_model_len: int
+    original_max_model_len: int | None
+    runner_type: str = "generate"
+    quantization: str | None = None
+    quantization_config: Any | None = None
+    hf_config: Any | None = None
+    hf_text_config: Any | None = None
+    multimodal_config: Any | None = None
+    enforce_eager: bool = False
+    disable_cascade_attn: bool = False
+    enable_return_routed_experts: bool = False
+    enable_prompt_embeds: bool = False
+    logits_processors: tuple[str, ...] = ()
+    use_mla: bool = False
+    is_moe: bool = False
+
+    @property
+    def is_quantized(self) -> bool:
+        return self.quantization is not None
+
+    def is_nvfp4_quantized(self) -> bool:
+        return self.quantization == "modelopt_fp4"
+
+    @property
+    def is_diffusion(self) -> bool:
+        # This config is also installed as the current VllmConfig for the
+        # existing Omni diffusion runtime, which does not use vLLM's native
+        # diffusion runner. Preserve that established backend-selection flag.
+        return False
+
+
+def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig:
+    quant_config = getattr(od_config, "quantization_config", None)
+    quantization = quant_config.get_name() if quant_config is not None and hasattr(quant_config, "get_name") else None
+    hf_config = getattr(od_config, "tf_model_config", None)
+    hf_text_config = get_hf_text_config(hf_config) if hasattr(hf_config, "get_text_config") else hf_config
+    configured_max_model_len = getattr(od_config, "max_model_len", None)
+    max_model_len = resolve_diffusion_max_model_len(od_config)
+    paged_scheduler = (
+        getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+        is DiffusionKVCacheMode.PAGED_SCHEDULER
+    )
+    return _DiffusionVllmModelConfig(
+        model=od_config.model,
+        dtype=od_config.dtype,
+        max_model_len=max_model_len,
+        # Preserve the user's -1 sentinel only for the native paged-cache
+        # sizing flow. Dense keeps its established resolved model-config view.
+        original_max_model_len=configured_max_model_len if paged_scheduler else max_model_len,
+        quantization=quantization,
+        quantization_config=quant_config,
+        hf_config=hf_config,
+        hf_text_config=hf_text_config,
+        enforce_eager=getattr(od_config, "enforce_eager", False),
+        is_moe=bool(getattr(od_config, "is_moe", False)),
+    )
+
+
+def create_base_diffusion_vllm_config(
+    device: torch.device,
+    od_config: OmniDiffusionConfig,
+) -> VllmConfig:
+    """Create the native vLLM 0.27 config used by diffusion."""
+
+    return VllmConfig(
+        compilation_config=CompilationConfig(),
+        device_config=DeviceConfig(device=device),
+        additional_config=od_config.additional_config,
+    )
+
+
+def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiffusionConfig) -> VllmConfig:
+    """Populate native model, parallel, cache, and Scheduler configuration."""
+
+    parallel_config = od_config.parallel_config
+    vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size
+    vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size
+    if parallel_config.enable_expert_parallel and od_config.is_moe:
+        vllm_config.parallel_config.data_parallel_size = (
+            parallel_config.data_parallel_size * parallel_config.cfg_parallel_size
+        )
+        vllm_config.parallel_config.prefill_context_parallel_size = parallel_config.sequence_parallel_size
+    vllm_config.parallel_config.enable_expert_parallel = parallel_config.enable_expert_parallel
+
+    vllm_config.model_config = _make_diffusion_vllm_model_config(od_config)  # type: ignore[assignment]
+    vllm_config.quant_config = od_config.quantization_config
+    vllm_config.profiler_config = od_config.profiler_config
+    if (
+        getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+        is DiffusionKVCacheMode.PAGED_SCHEDULER
+    ):
+        vllm_config.parallel_config.pipeline_parallel_size = parallel_config.pipeline_parallel_size
+        vllm_config.cache_config.gpu_memory_utilization = float(getattr(od_config, "gpu_memory_utilization", 0.9))
+        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
+        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)
+    return vllm_config
+
+
+def create_diffusion_vllm_config(
+    device: torch.device,
+    od_config: OmniDiffusionConfig,
+) -> VllmConfig:
+    """Build the same fully populated native config in Engine and Workers."""
+
+    vllm_config = create_base_diffusion_vllm_config(
+        device,
+        od_config,
+    )
+    return configure_diffusion_vllm_config(vllm_config, od_config)

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -11,6 +11,7 @@ model-related operations.
 from __future__ import annotations
 
 import copy
+import gc
 import time
 from collections.abc import Callable
 from contextlib import AbstractContextManager, nullcontext
@@ -21,6 +22,7 @@ from torch.profiler import record_function
 from vllm.config import LoadConfig, VllmConfig
 from vllm.logger import init_logger
 from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 
 from vllm_omni.diffusion.cache.cachedit import cache_summary
 from vllm_omni.diffusion.cache.prompt_embed_cache import (
@@ -45,8 +47,10 @@ from vllm_omni.diffusion.offloader import get_offload_backend
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
+    CachedRequestData,
     DiffusionSchedulerOutput,
     KVPrefetchJob,
+    NewRequestData,
     validate_new_request_data_identity,
 )
 from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
@@ -139,6 +143,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         self.offload_backend: Any | None = None
         self.prompt_embed_cache: Any | None = None
         self.input_batch: InputBatch | None = None
+        self.model_memory_usage = 0
+        # Rank-local native config generated by the Engine. Physical tensors
+        # and BlockTable binding remain the responsibility of the paged
+        # data-plane initialization path.
+        self.kv_cache_config: KVCacheConfig | None = None
 
         # Cache for per-request stepwise state.
         self.state_cache: dict[str, StepRequestState] = {}
@@ -167,8 +176,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         metadata: DiffusionKVMetadata | None,
     ) -> None:
         cache_mode = getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
-        # W0 does not wire the Scheduler allocator or Worker installer yet, so
-        # validate an allocation snapshot when present without requiring one.
+        if cache_mode is DiffusionKVCacheMode.PAGED_SCHEDULER and metadata is None:
+            raise ValueError(f"paged_scheduler request {request_id!r} requires Diffusion KV metadata")
         if cache_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER and metadata is not None:
             raise ValueError(f"{cache_mode.value} request {request_id!r} must not carry Diffusion KV metadata")
 
@@ -273,6 +282,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             m.consumed_memory / GiB_bytes,
             time_after_load - time_before_load,
         )
+        self.model_memory_usage = int(m.consumed_memory)
         logger.info("Model runner: Model loaded successfully.")
 
         if self.od_config.streaming_output and not getattr(self.od_config, "step_execution", False):
@@ -368,6 +378,45 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
         logger.info("Model runner: Initialization complete.")
 
+    def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
+        """Collect native specs from cache-enabled loaded attention modules."""
+
+        from vllm_omni.diffusion.attention.layer import Attention
+        from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+
+        if self.od_config.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
+            return {}
+        if self.pipeline is None:
+            raise RuntimeError("Model must be loaded before collecting Diffusion KV cache specs")
+
+        specs: dict[str, KVCacheSpec] = {}
+        for layer_name, module in self.pipeline.named_modules():
+            if not isinstance(module, Attention):
+                continue
+            spec = module.get_kv_cache_spec(self.vllm_config)
+            if spec is not None:
+                specs[layer_name] = spec
+        if not specs:
+            raise RuntimeError(
+                "paged_scheduler Diffusion KV found no cache-enabled Attention modules "
+                f"in {type(self.pipeline).__name__}"
+            )
+        return specs
+
+    def set_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None:
+        """Retain the Engine-generated rank-local config for the data plane."""
+
+        configured_layers = {
+            layer_name for group in kv_cache_config.kv_cache_groups for layer_name in group.layer_names
+        }
+        expected_layers = set(self.get_kv_cache_spec())
+        if configured_layers != expected_layers:
+            raise ValueError(
+                "Rank-local Diffusion KVCacheConfig layer mismatch: "
+                f"expected={sorted(expected_layers)}, configured={sorted(configured_layers)}"
+            )
+        self.kv_cache_config = kv_cache_config
+
     def clear_prompt_embed_cache(self) -> None:
         """Evict all cached text-encoder outputs (e.g. between training epochs).
 
@@ -516,6 +565,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         require_request_batch_support: bool,
         kv_prefetch_job: KVPrefetchJob | None = None,
         record_name: str,
+        record_output_peak_memory: bool = True,
     ) -> BatchRunnerOutput:
         assert self.pipeline is not None, "Model not loaded. Call load_model() first."
         if not reqs:
@@ -545,7 +595,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
             batch = DiffusionRequestBatch(requests=reqs)
             is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
-            if is_primary:
+            if is_primary and record_output_peak_memory:
                 current_omni_platform.reset_peak_memory_stats()
 
             with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=od_config):
@@ -558,7 +608,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                         pipeline_name=type(self.pipeline).__name__,
                     )
 
-            if is_primary and outputs:
+            if is_primary and outputs and record_output_peak_memory:
                 batch_peak_memory_mb = self._sample_peak_memory_mb()
                 for output in outputs:
                     output.peak_memory_mb = max(output.peak_memory_mb, batch_peak_memory_mb)
@@ -626,6 +676,59 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         assert output is not None
         return output
 
+    def profile_run(self, requests: list[OmniDiffusionRequest]) -> None:
+        """Run the maximum per-rank request batch for memory profiling.
+
+        This deliberately bypasses Scheduler admission and Diffusion KV
+        metadata validation because cache capacity has not been sized yet.
+        It otherwise uses the normal execution-mode path so model inputs,
+        collective communication, backend workspaces, denoising, and decode
+        allocations contribute to the observed peak. Step execution profiles
+        one fused ``InputBatch`` instead of sequential single-request forwards.
+        """
+
+        if not requests:
+            raise ValueError("Diffusion memory profiling requires at least one request.")
+
+        runner_output: BatchRunnerOutput | None = None
+        request_ids = [request.request_id for request in requests]
+        try:
+            if getattr(self.od_config, "step_execution", False):
+                scheduler_output = DiffusionSchedulerOutput(
+                    step_id=0,
+                    scheduled_new_reqs=[
+                        NewRequestData(request_id=request.request_id, req=request) for request in requests
+                    ],
+                    scheduled_cached_reqs=CachedRequestData.make_empty(),
+                    finished_req_ids=set(),
+                    num_running_reqs=len(requests),
+                    num_waiting_reqs=0,
+                )
+                runner_output = self._execute_stepwise(
+                    scheduler_output,
+                    validate_kv_metadata=False,
+                    record_output_peak_memory=False,
+                )
+            else:
+                runner_output = self._execute_request_list(
+                    requests,
+                    od_config=self.od_config,
+                    allow_single_output=len(requests) == 1,
+                    require_request_batch_support=len(requests) > 1,
+                    record_name="pipeline_memory_profile",
+                    # The enclosing native memory_profiling context owns the
+                    # peak counters. Resetting them here would discard request
+                    # preparation allocations and understate the budget.
+                    record_output_peak_memory=False,
+                )
+            current_omni_platform.synchronize()
+        finally:
+            for request_id in request_ids:
+                self.state_cache.pop(request_id, None)
+            self.input_batch = None
+            del runner_output
+            gc.collect()
+
     def execute_model_batch(
         self,
         scheduler_output: DiffusionSchedulerOutput,
@@ -750,13 +853,29 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
     def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput:
         """Execute one step for one scheduled request and return runner output."""
+        return self._execute_stepwise(
+            scheduler_output,
+            validate_kv_metadata=True,
+            record_output_peak_memory=True,
+        )
+
+    def _execute_stepwise(
+        self,
+        scheduler_output: DiffusionSchedulerOutput,
+        *,
+        validate_kv_metadata: bool,
+        record_output_peak_memory: bool,
+    ) -> BatchRunnerOutput:
+        """Execute one step with explicit validation and profiling policy."""
+
         assert self.pipeline is not None, "Model not loaded. Call load_model() first."
         for new_req in scheduler_output.scheduled_new_reqs:
             validate_new_request_data_identity(new_req)
-            self._validate_diffusion_kv_metadata(
-                request_id=new_req.req.request_id,
-                metadata=new_req.diffusion_kv_metadata,
-            )
+            if validate_kv_metadata:
+                self._validate_diffusion_kv_metadata(
+                    request_id=new_req.req.request_id,
+                    metadata=new_req.diffusion_kv_metadata,
+                )
         if not self._supports_step_mode():
             raise ValueError("Current pipeline does not support step execution.")
         # Stepwise mode only supports the basic state-driven denoise path for now.
@@ -771,7 +890,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             had_active_states = bool(self.state_cache)
             states, new_request_ids = self._update_states(scheduler_output)
             is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
-            if new_request_ids and not had_active_states and is_primary and current_omni_platform.is_available():
+            if (
+                record_output_peak_memory
+                and new_request_ids
+                and not had_active_states
+                and is_primary
+                and current_omni_platform.is_available()
+            ):
                 current_omni_platform.reset_peak_memory_stats()
             input_batch = self._prepare_batch_inputs(states, new_request_ids)
             attn_metadata = {}
@@ -864,7 +989,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                             f"but batched noise_pred has {noise_pred.shape[0]} rows."
                         )
 
-                if is_primary:
+                if is_primary and record_output_peak_memory:
                     batch_peak_memory_mb = self._sample_peak_memory_mb()
                     states_by_id = {state.request_id: state for state in states}
                     for state in states:

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -16,20 +16,21 @@ import signal
 import threading
 import traceback
 import uuid
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
-from dataclasses import dataclass
 from typing import Any
 
 import torch
+import torch.distributed as dist
 import zmq
-from vllm.config import CompilationConfig, DeviceConfig, VllmConfig, set_current_vllm_config
+from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.logger import init_logger
 from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler
-from vllm.transformers_utils.config import get_hf_text_config
 from vllm.utils.import_utils import resolve_obj_by_qualname
-from vllm.utils.mem_utils import GiB_bytes
+from vllm.utils.mem_utils import GiB_bytes, MemorySnapshot, format_gib, memory_profiling
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
+from vllm.v1.worker.utils import request_memory
 from vllm.v1.worker.workspace import init_workspace_manager
 
 from vllm_omni.diffusion.data import (
@@ -41,6 +42,7 @@ from vllm_omni.diffusion.data import (
     OmniSleepTask,
     OmniWakeTask,
 )
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -52,7 +54,13 @@ from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusio
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager, LoRABackend
 from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
+from vllm_omni.diffusion.sched.interface import (
+    DiffusionSchedulerOutput,
+    KVPrefetchJob,
+    NewRequestData,
+    validate_new_request_data_identity,
+)
+from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput
 from vllm_omni.engine.stage_init_utils import set_death_signal
@@ -76,49 +84,28 @@ _ASYNC_OUTPUT_DRAIN_TIMEOUT_S = 10.0
 _MEMORY_RELEASING_METHODS = frozenset({"sleep", "handle_sleep_task"})
 
 
-@dataclass
-class _DiffusionVllmModelConfig:
-    model: str
-    dtype: torch.dtype
-    quantization: str | None = None
-    quantization_config: Any | None = None
-    hf_config: Any | None = None
-    hf_text_config: Any | None = None
-    multimodal_config: Any | None = None
-    enforce_eager: bool = False
-    disable_cascade_attn: bool = False
-    enable_return_routed_experts: bool = False
-    is_moe: bool = False
-
-    def is_quantized(self) -> bool:
-        return self.quantization is not None
-
-    def is_model_moe(self) -> bool:
-        return self.is_moe
-
-    def is_nvfp4_quantized(self) -> bool:
-        return self.quantization == "modelopt_fp4"
-
-    @property
-    def is_diffusion(self) -> bool:
-        return False
+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)
+    return values
 
 
-def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig:
-    quant_config = getattr(od_config, "quantization_config", None)
-    quantization = quant_config.get_name() if quant_config is not None and hasattr(quant_config, "get_name") else None
-    hf_config = getattr(od_config, "tf_model_config", None)
-    hf_text_config = get_hf_text_config(hf_config) if hasattr(hf_config, "get_text_config") else hf_config
-    return _DiffusionVllmModelConfig(
-        model=od_config.model,
-        dtype=od_config.dtype,
-        quantization=quantization,
-        quantization_config=quant_config,
-        hf_config=hf_config,
-        hf_text_config=hf_text_config,
-        enforce_eager=getattr(od_config, "enforce_eager", False),
-        is_moe=bool(getattr(od_config, "is_moe", False)),
-    )
+def _run_and_gather_rank_values(operation: str, func: Callable[[], Any]) -> list[Any]:
+    """Run one rank-local probe without stranding peers on local failure."""
+
+    try:
+        local_result = (True, func())
+    except Exception as exc:
+        logger.exception("%s failed on this Worker rank", operation)
+        local_result = (False, f"{type(exc).__name__}: {exc}")
+
+    rank_results = _all_gather_rank_values(local_result)
+    failures = [f"rank {rank}: {result}" for rank, (ok, result) in enumerate(rank_results) if not ok]
+    if failures:
+        raise RuntimeError(f"{operation} failed on " + "; ".join(failures))
+    return [result for _, result in rank_results]
 
 
 @contextmanager
@@ -152,37 +139,6 @@ def _force_cutlass_fp8_linear_kernel(quant_config: object | None) -> Iterator[No
     yield
 
 
-def _is_unexpected_additional_config_type_error(exc: TypeError) -> bool:
-    """Return True only for constructor rejections of the additional_config kwarg."""
-    message = str(exc)
-    return "unexpected keyword argument" in message and "additional_config" in message
-
-
-def _create_diffusion_worker_vllm_config(device: torch.device, od_config: OmniDiffusionConfig) -> VllmConfig:
-    """Create a worker-local VllmConfig while preserving additional_config when supported."""
-    config_kwargs: dict[str, Any] = {
-        "compilation_config": CompilationConfig(),
-        "device_config": DeviceConfig(device=device),
-    }
-    if od_config.additional_config:
-        config_kwargs["additional_config"] = od_config.additional_config
-
-    try:
-        return VllmConfig(**config_kwargs)
-    except TypeError as exc:
-        if not _is_unexpected_additional_config_type_error(exc):
-            raise
-
-        logger.debug("Worker-local VllmConfig does not accept additional_config in constructor: %s", exc)
-        config_kwargs.pop("additional_config", None)
-        vllm_config = VllmConfig(**config_kwargs)
-        try:
-            setattr(vllm_config, "additional_config", dict(od_config.additional_config))
-        except Exception as set_exc:  # pragma: no cover - defensive for older vLLM builds
-            logger.warning("Failed to attach additional_config to worker VllmConfig: %s", set_exc)
-        return vllm_config
-
-
 def _get_cumem_allocator_class() -> type:
     from vllm.device_allocator.cumem import CuMemAllocator
 
@@ -223,6 +179,8 @@ class DiffusionWorker:
         self.device: torch.device | None = None
         self.vllm_config: VllmConfig | None = None
         self.model_runner: DiffusionModelRunner | None = None
+        self.init_snapshot: MemorySnapshot | None = None
+        self.requested_memory: int | None = None
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
         self.lora_manager: DiffusionLoRAManager | None = None
         # Worker-side cache of (lora_request, lora_scale) per scheduled
@@ -287,22 +245,7 @@ class DiffusionWorker:
 
         # Create vllm_config for parallel configuration. Pass explicit device_config
         # so DeviceConfig does not rely on current_platform in worker subprocesses.
-        vllm_config = _create_diffusion_worker_vllm_config(self.device, self.od_config)
-        parallel_config = self.od_config.parallel_config
-        vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size
-        vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size
-        if parallel_config.enable_expert_parallel and self.od_config.is_moe:
-            # Diffusion uses its own DP/CFG/SP groups normally. vLLM groups are
-            # only remapped for expert-parallel runtimes that consume vLLM's
-            # FusedMoE/EP semantics.
-            vllm_config.parallel_config.data_parallel_size = (
-                parallel_config.data_parallel_size * parallel_config.cfg_parallel_size
-            )
-            vllm_config.parallel_config.prefill_context_parallel_size = parallel_config.sequence_parallel_size
-        vllm_config.parallel_config.enable_expert_parallel = parallel_config.enable_expert_parallel
-        vllm_config.profiler_config = self.od_config.profiler_config
-        vllm_config.model_config = _make_diffusion_vllm_model_config(self.od_config)  # type: ignore[assignment]
-        vllm_config.quant_config = self.od_config.quantization_config
+        vllm_config = create_diffusion_vllm_config(self.device, self.od_config)
         # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping.
         # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion.
         vllm_config.kernel_config.ir_op_priority = _resolve_ir_op_priority(self.od_config, vllm_config)
@@ -340,6 +283,23 @@ class DiffusionWorker:
                 hsdp_replicate_size=parallel_config.hsdp_replicate_size if parallel_config.use_hsdp else 1,
                 enable_expert_parallel=parallel_config.enable_expert_parallel,
             )
+            if (
+                getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+                is DiffusionKVCacheMode.PAGED_SCHEDULER
+            ):
+                gc.collect()
+                current_omni_platform.empty_cache()
+                self.init_snapshot = MemorySnapshot(device=self.device)
+                self.requested_memory = request_memory(
+                    self.init_snapshot,
+                    vllm_config.cache_config,
+                )
+                logger.debug(
+                    "Worker %d: Diffusion KV initial memory snapshot: %r; requested=%s GiB",
+                    self.rank,
+                    self.init_snapshot,
+                    format_gib(self.requested_memory),
+                )
             init_workspace_manager(self.device)
 
     def _create_profiler(self) -> WorkerProfiler | None:
@@ -393,6 +353,96 @@ class DiffusionWorker:
         if load_format != "dummy":
             assert self.model_runner.pipeline is not None
 
+    def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
+        """Return native rank-local specs for every diffusion Worker."""
+
+        assert self.model_runner is not None
+        return _run_and_gather_rank_values(
+            "Diffusion KV cache spec discovery",
+            self.model_runner.get_kv_cache_spec,
+        )
+
+    def determine_available_kv_memory(self, profile_requests: list[OmniDiffusionRequest]) -> list[int]:
+        """Profile and return each rank's safe Diffusion KV memory budget."""
+
+        def determine_local_memory() -> int:
+            assert self.vllm_config is not None
+            assert self.model_runner is not None
+            if self.init_snapshot is None or self.requested_memory is None:
+                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)
+                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.",
+                    self.rank,
+                    format_gib(self.init_snapshot.free_memory),
+                    format_gib(int(override)),
+                )
+                return int(override)
+
+            with memory_profiling(
+                self.init_snapshot,
+                weights_memory=self.model_runner.model_memory_usage,
+            ) as profile_result:
+                self.model_runner.profile_run(profile_requests)
+
+            available_memory = self.requested_memory - profile_result.non_kv_cache_memory
+            if available_memory <= 0:
+                raise RuntimeError(
+                    "No memory remains for Diffusion KV cache after profiling: "
+                    f"requested_memory={self.requested_memory} bytes, "
+                    f"non_kv_cache_memory={profile_result.non_kv_cache_memory} bytes. "
+                    "Increase gpu_memory_utilization or reduce the maximum profile request shape."
+                )
+            free_gpu_memory = profile_result.after_profile.free_memory
+            unrequested_memory = self.init_snapshot.free_memory - self.requested_memory
+            logger.debug(
+                "Worker %d: Initial free memory: %s GiB; Requested memory: %f (util), %s GiB",
+                self.rank,
+                format_gib(self.init_snapshot.free_memory),
+                self.vllm_config.cache_config.gpu_memory_utilization,
+                format_gib(self.requested_memory),
+            )
+            logger.debug(
+                "Worker %d: Free memory after profiling: %s GiB (total), %s GiB (within requested)",
+                self.rank,
+                format_gib(free_gpu_memory),
+                format_gib(free_gpu_memory - unrequested_memory),
+            )
+            logger.debug("Worker %d: %r", self.rank, profile_result)
+            logger.info(
+                "Worker %d: Available Diffusion KV cache memory: %s GiB",
+                self.rank,
+                format_gib(available_memory),
+            )
+            return int(available_memory)
+
+        return [
+            int(value)
+            for value in _run_and_gather_rank_values(
+                "Diffusion KV memory discovery",
+                determine_local_memory,
+            )
+        ]
+
+    def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig]) -> None:
+        """Install this rank's native config without allocating tensors yet."""
+
+        assert self.model_runner is not None
+        if len(kv_cache_configs) != self.od_config.num_gpus:
+            raise ValueError(
+                "Diffusion KVCacheConfig rank count mismatch: "
+                f"expected={self.od_config.num_gpus}, got={len(kv_cache_configs)}"
+            )
+        self.model_runner.set_kv_cache_config(kv_cache_configs[self.rank])
+
     def init_lora_manager(self) -> None:
         """Initialize the LoRA manager for this worker."""
         if self.model_runner.pipeline is None:
@@ -465,18 +515,19 @@ class DiffusionWorker:
 
     def execute_model(
         self,
-        req: OmniDiffusionRequest | list[OmniDiffusionRequest],
+        req: OmniDiffusionRequest | list[NewRequestData],
         od_config: OmniDiffusionConfig,
         kv_prefetch_job: KVPrefetchJob | None = None,
         diffusion_kv_metadata: DiffusionKVMetadata | None = None,
     ) -> DiffusionOutput:
         """Execute a forward pass by delegating to the model runner.
 
-        If *req* is a list (DP multi-concurrency), each rank picks one
-        request based on its distributed rank.  AllGather in the layerwise
-        offload only gathers weight shards (request-independent), so all
-        ranks stay synchronised at each AllGather call while computing
-        different activations.
+        If *req* is a list (DP multi-concurrency), each rank picks one complete
+        NewRequestData envelope based on its distributed rank. AllGather in
+        the layerwise offload only gathers weight shards (request-independent),
+        so all ranks stay synchronised at each AllGather call while computing
+        different activations. Selecting the envelope keeps Scheduler-issued
+        KV metadata bound to the request that owns its block tables.
 
         Each rank returns its OWN DiffusionOutput (no gather). The executor
         collects N responses via the per-worker result queues.
@@ -492,7 +543,10 @@ class DiffusionWorker:
 
             dp_rank = get_data_parallel_rank()
             idx = dp_rank % len(req)
-            req = req[idx]
+            new_req = req[idx]
+            validate_new_request_data_identity(new_req)
+            req = new_req.req
+            diffusion_kv_metadata = new_req.diffusion_kv_metadata
 
         if self.lora_manager is not None:
             try:
@@ -1443,7 +1497,7 @@ class WorkerWrapperBase:
 
     def execute_model(
         self,
-        req: OmniDiffusionRequest,
+        req: OmniDiffusionRequest | list[NewRequestData],
         od_config: OmniDiffusionConfig,
         kv_prefetch_job: KVPrefetchJob | None = None,
         diffusion_kv_metadata: DiffusionKVMetadata | None = None,