[Diffusion] Register runtime cache for direct DLO H2D

Avoid recurrent mmap-to-pinned packing in no-AllGather DLO by registering transform-complete runtime-cache mappings under an explicit per-worker budget. Copy mapped tensors directly into the rotating device buffers while preserving bounded-staging fallback and ordered shutdown cleanup. Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:21:28 +08:00
1 parent 67f15e73cd
commit 23ab14a15b
@@ -492,6 +492,59 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch,
    assert loader.take_host_weight_plan() is runtime_plan


+def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=1,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        dlo_runtime_cache_pin_limit_gib=1.0,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+    runtime_plan = HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings={},
+        runtime_layout_key="layout-key",
+        post_load_complete=True,
+    )
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: events.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: events.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+
+    def unexpected_checkpoint_plan(*_args, **_kwargs):
+        pytest.fail("registered runtime-cache mode must not select the raw checkpoint mmap plan")
+
+    def build_cache(_pipeline, **_kwargs):
+        events.append("cache")
+        return HostWeightPlanResult(runtime_plan)
+
+    monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan)
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert events == ["load", "process", "cache"]
+    assert loader.take_host_weight_plan() is runtime_plan
+
+
 def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
+from vllm_omni.diffusion.offloader.cuda_host_registration import (
+    CudaHostRegistration,
+    CudaHostRegistrationBudgetError,
+    CudaHostRegistrationCleanupError,
+    CudaHostRegistrationError,
+    _coalesce_ranges,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _FakeStorage:
+    def __init__(self, pointer: int, size: int) -> None:
+        self._pointer = pointer
+        self._size = size
+
+    def data_ptr(self) -> int:
+        return self._pointer
+
+    def nbytes(self) -> int:
+        return self._size
+
+
+class _FakeTensor:
+    device = SimpleNamespace(type="cpu")
+
+    def __init__(self, pointer: int, size: int, *, pinned: bool = True) -> None:
+        self._storage = _FakeStorage(pointer, size)
+        self._pinned = pinned
+
+    def untyped_storage(self) -> _FakeStorage:
+        return self._storage
+
+    def numel(self) -> int:
+        return self._storage.nbytes()
+
+    def is_pinned(self) -> bool:
+        return self._pinned
+
+
+class _FakeRuntime:
+    def __init__(
+        self,
+        register_results: list[int | Exception],
+        unregister_results: list[int | Exception] | None = None,
+    ) -> None:
+        self._register_results = iter(register_results)
+        self._unregister_results = iter(unregister_results or [])
+        self.registered: list[tuple[int, int, int]] = []
+        self.unregistered: list[int] = []
+
+    def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
+        self.registered.append((pointer, size, flags))
+        result = next(self._register_results)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    def cudaHostUnregister(self, pointer: int) -> int:
+        self.unregistered.append(pointer)
+        result = next(self._unregister_results, 0)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    @staticmethod
+    def cudaGetErrorString(error: int) -> str:
+        return f"error-{error}"
+
+
+def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
+    assert _coalesce_ranges(
+        [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
+        page_size=4096,
+    ) == (
+        registration_module._AddressRange(0x1000, 0x3000),
+        registration_module._AddressRange(0x9000, 0xA000),
+    )
+
+
+def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationBudgetError, match="exceeding"):
+        CudaHostRegistration.create(
+            {"weights": [_FakeTensor(0x1003, 4096)]},  # type: ignore[list-item]
+            max_bytes=4096,
+        )
+
+    assert runtime.registered == []
+
+
+def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="error-7"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_wraps_runtime_exception_and_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, RuntimeError("driver rejected mapping")])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="driver rejected mapping"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7], unregister_results=[9])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationCleanupError, match="rollback errors"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+
+def test_registration_closes_all_coalesced_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {
+            "first": [_FakeTensor(0x1003, 4096), _FakeTensor(0x2800, 1024)],  # type: ignore[list-item]
+            "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+        },
+        max_bytes=12288,
+    )
+
+    assert registration.total_bytes == 12288
+    assert registration.region_count == 2
+    assert registration.close() == []
+    assert runtime.unregistered == [0x9000, 0x1000]
+    assert registration.close() == []
+
+
+def test_registration_keeps_adjacent_file_mappings_separate(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    registration = CudaHostRegistration.create(
+        {
+            "first.safetensors": [_FakeTensor(0x1000, 4096)],  # type: ignore[list-item]
+            "second.safetensors": [_FakeTensor(0x2000, 4096)],  # type: ignore[list-item]
+        },
+        max_bytes=8192,
+    )
+
+    assert registration.region_count == 2
+    assert runtime.registered == [
+        (0x1000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+        (0x2000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+    ]
+    assert registration.close() == []
+
+
+def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0], unregister_results=[9, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
+        max_bytes=4096,
+    )
+
+    assert registration.close() == ["cudaHostUnregister(0x1000) failed: error-9"]
+    assert registration.close() == []
+    assert runtime.unregistered == [0x1000, 0x1000]
@@ -546,6 +546,31 @@ class TestDistributedLayerwiseOffloadHook:
        assert staging[0][torch.float32].numel() == 12
        assert staging[1][torch.float32].numel() == 12

+    def test_registered_mmap_copies_directly_without_staging(self, patched_offload_runtime):
+        current_block = nn.Linear(2, 2, bias=False)
+        next_block = nn.Linear(2, 2, bias=False)
+        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+        next_block.weight.data.copy_(expected)
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        hook.initialize_hook(current_block)
+        hook.registered_mmap = True
+
+        def fail_staging(_slot):
+            raise AssertionError("registered mmap must bypass host staging")
+
+        hook._stage_mmap_sources = fail_staging  # type: ignore[method-assign]
+        hook.prefetch_layer(slot=0, non_blocking=True)
+
+        assert torch.equal(next_block.weight, expected)
+

 class TestPinnedResidentLayerGroup:
    def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -636,6 +661,24 @@ class TestPinnedResidentLayerGroup:
        group.offload()
        assert all(block.weight.numel() == 0 for block in blocks)

+    def test_registered_mmap_resident_layers_bypass_staging(self, patched_offload_runtime):
+        block = nn.Linear(2, 2, bias=False)
+        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+        block.weight.data.copy_(expected)
+        group = PinnedResidentLayerGroup(
+            [block],
+            device=torch.device("cpu"),
+            copy_stream=DummyStream(),
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        group.registered_mmap = True
+        group._cpu_staging_buffers.clear()
+
+        group.load()
+
+        assert torch.equal(block.weight, expected)
+
    def test_shared_allgather_output_is_narrowed_to_current_block(
        self, dist_group, patched_offload_runtime, monkeypatch
    ):
@@ -995,6 +1038,52 @@ class TestMmapWeightLoading:
        assert len(cleanup_calls) == 2
        backend.disable()

+    def test_runtime_cache_registration_uses_budget_and_releases_before_mmap(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        patched_offload_runtime,
+    ):
+        events: list[str] = []
+
+        class Registration:
+            total_bytes = 4096
+            region_count = 1
+
+            @staticmethod
+            def close() -> list[str]:
+                events.append("unregister")
+                return []
+
+        create_calls: list[tuple[dict[str, list[torch.Tensor]], int]] = []
+
+        def create(sources, *, max_bytes):
+            create_calls.append((sources, max_bytes))
+            return Registration()
+
+        monkeypatch.setattr(dist_backend_module.CudaHostRegistration, "create", staticmethod(create))
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dlo_use_allgather=False,
+                dlo_runtime_cache_pin_limit_gib=1.5,
+            ),
+            torch.device("cuda"),
+        )
+        sources = {"runtime.safetensors": [torch.ones(1)]}
+        backend._runtime_cache_mapped_sources = sources
+
+        assert backend._try_register_runtime_cache_mmap()
+        assert create_calls == [(sources, int(1.5 * 1024**3))]
+
+        backend.enabled = True
+        backend._using_rank_local_mmap = True
+        backend._mmap_file_cache = {}
+        backend._release_mmap_handles = lambda: events.append("mmap")  # type: ignore[method-assign]
+        backend.disable()
+
+        assert events == ["unregister", "mmap"]
+

 class TestGetBlocksFromDit:
    def test_get_blocks_from_dit_single_block_attr(self):
@@ -12,11 +12,13 @@ This module tests the WorkerWrapperBase implementation:
 - Dynamic worker class extension
 """

+from types import SimpleNamespace
 from typing import Any

 import pytest
 from pytest_mock import MockerFixture

+import vllm_omni.diffusion.worker.diffusion_worker as worker_module
 from vllm_omni.diffusion.worker.diffusion_worker import (
    CustomPipelineWorkerExtension,
    DiffusionWorker,
@@ -204,6 +206,29 @@ class TestWorkerWrapperBaseDelegation:
        wrapper.worker.shutdown.assert_called_once()
        assert result is None

+    def test_worker_shutdown_disables_offloader_before_distributed_teardown(
+        self,
+        mocker: MockerFixture,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        worker = object.__new__(DiffusionWorker)
+        offload_backend = mocker.Mock()
+        kv_transfer_manager = mocker.Mock()
+        worker.model_runner = SimpleNamespace(
+            offload_backend=offload_backend,
+            kv_transfer_manager=kv_transfer_manager,
+        )
+        events: list[str] = []
+        offload_backend.disable.side_effect = lambda: events.append("offloader")
+        kv_transfer_manager.shutdown_prefetch.side_effect = lambda: events.append("kv_transfer")
+        monkeypatch.setattr(worker_module, "destroy_distributed_env", lambda: events.append("distributed"))
+
+        worker.shutdown()
+
+        offload_backend.disable.assert_called_once_with()
+        kv_transfer_manager.shutdown_prefetch.assert_called_once_with()
+        assert events == ["offloader", "kv_transfer", "distributed"]
+

 # -------------------------------------------------------------------------
 # Tests: execute_method