[Diffusion] Register HWR mmap for direct DLO H2D (#6591)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>

tests/diffusion/test_worker_wrapper_base.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Unit tests for WorkerWrapperBase class.
@@ -12,6 +12,7 @@ This module tests the WorkerWrapperBase implementation:
 - Dynamic worker class extension
 """
 
+from types import SimpleNamespace
 from typing import Any
 
 import pytest
@@ -204,6 +205,27 @@ class TestWorkerWrapperBaseDelegation:
         wrapper.worker.shutdown.assert_called_once()
         assert result is None
 
+    def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture):
+        events: list[str] = []
+        offload_backend = mocker.Mock()
+        offload_backend.disable.side_effect = lambda: events.append("offload")
+        kv_manager = mocker.Mock()
+        kv_manager.shutdown_prefetch.side_effect = lambda: events.append("kv")
+        destroy = mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env",
+            side_effect=lambda: events.append("distributed"),
+        )
+        worker = DiffusionWorker.__new__(DiffusionWorker)
+        worker.model_runner = SimpleNamespace(
+            offload_backend=offload_backend,
+            kv_transfer_manager=kv_manager,
+        )
+
+        worker.shutdown()
+
+        assert events == ["offload", "kv", "distributed"]
+        destroy.assert_called_once_with()
+
 
 # -------------------------------------------------------------------------
 # Tests: execute_method