[Diffusion] Add paged KV cache support for HunyuanImage3 DiT on NPU and GPU (#6563)

Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Signed-off-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: zwhzzz0821 <2831474076@qq.com>

tests/platforms/npu/test_diffusion_platform.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import importlib.util
+from types import SimpleNamespace
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("available", [False, True])
+def test_dense_flash_capability_tracks_mindiesd(monkeypatch, available):
+    pytest.importorskip("vllm_ascend")
+    from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+    monkeypatch.setattr(
+        importlib.util,
+        "find_spec",
+        lambda name: object() if name == "mindiesd" and available else None,
+    )
+
+    assert NPUOmniPlatform.supports_diffusion_dense_flash_attention() is available
+
+
+def test_paged_config_uses_ascend_kernel_block_size() -> None:
+    pytest.importorskip("vllm_ascend")
+    from vllm_ascend.attention.attention_v1 import AscendAttentionBackend
+
+    from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+    from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+    vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16))
+    NPUOmniPlatform.configure_diffusion_vllm_config(
+        vllm_config,
+        SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER),
+    )
+
+    assert vllm_config.cache_config.block_size == AscendAttentionBackend.get_supported_kernel_block_sizes()[0]
+
+
+def test_strict_ulysses_paged_backend_bypasses_ascend_pcp() -> None:
+    pytest.importorskip("vllm_ascend")
+    from vllm_ascend.attention.attention_v1 import (
+        AscendAttentionBackend,
+        AscendAttentionBackendImpl,
+        AscendAttentionMetadataBuilder,
+    )
+
+    from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+    assert (
+        NPUOmniPlatform.get_diffusion_paged_kv_attn_backend(
+            AscendAttentionBackend,
+            ulysses_degree=1,
+        )
+        is AscendAttentionBackend
+    )
+    backend = NPUOmniPlatform.get_diffusion_paged_kv_attn_backend(
+        AscendAttentionBackend,
+        ulysses_degree=2,
+    )
+    assert backend.get_impl_cls() is AscendAttentionBackendImpl
+    assert backend.get_builder_cls() is AscendAttentionMetadataBuilder

tests/platforms/test_record_device_event.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 OmniPlatform.record_device_event implementations."""
 
 import pytest
@@ -122,8 +122,8 @@ class TestNPUOmniPlatformRecordDeviceEvent:
         result = NPUOmniPlatform.record_device_event()
         assert result is None
 
-    def test_synchronizes_stream_then_records_event(self, mocker):
-        """NPU should sync stream first, then create and record an Event."""
+    def test_synchronizes_stream_then_records_generic_event(self, mocker):
+        """NPU should return an event consumable by the generic torch.Stream."""
         try:
             from vllm_omni.platforms.npu.platform import NPUOmniPlatform
         except ModuleNotFoundError:
@@ -133,14 +133,16 @@ class TestNPUOmniPlatformRecordDeviceEvent:
         mock_stream = mocker.MagicMock()
         mock_torch = mocker.MagicMock()
         mock_torch.npu.current_stream.return_value = mock_stream
-        mock_torch.npu.Event.return_value = mock_event
+        mock_torch.Event.return_value = mock_event
         mocker.patch("vllm_omni.platforms.npu.platform.torch", mock_torch)
 
         result = NPUOmniPlatform.record_device_event()
 
         # Stream should be synced first (HCCL ordering)
         mock_stream.synchronize.assert_called_once()
-        # Then event should be created and recorded
-        mock_torch.npu.Event.assert_called_once()
+        # Then a backend-neutral event should be created and recorded.  The
+        # worker waits on it through the public torch.Stream wrapper.
+        mock_torch.Event.assert_called_once()
+        mock_torch.npu.Event.assert_not_called()
         mock_event.record.assert_called_once()
         assert result is mock_event