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

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