[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/diffusion/attention/test_attention_config.py

@@ -11,6 +11,7 @@ Tests cover:
 """
 
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -466,6 +467,120 @@ class TestCurrentDiffusionConfig:
 
 
 class TestAttentionInitUsesCurrentDiffusionConfig:
+    @pytest.mark.parametrize("mindiesd_available", [False, True])
+    def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypatch, mindiesd_available):
+        attention = Attention.__new__(Attention)
+        attention._scheduler_paged_kv = True
+        attention.paged_kv_cache_role = "primary"
+        attention.backend_pref = "FLASH_ATTN"
+        attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
+        flash_output = torch.ones(1)
+        sdpa_output = torch.zeros(1)
+        dense_flash = Mock(return_value=flash_output)
+        sdpa = Mock(return_value=sdpa_output)
+        attention.attention = SimpleNamespace(forward=dense_flash)
+        attention.sdpa_fallback = SimpleNamespace(forward=sdpa)
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True)
+        monkeypatch.setattr(
+            layer_mod,
+            "get_forward_context",
+            lambda: SimpleNamespace(in_diffusion_kv_memory_profile=True),
+        )
+        monkeypatch.setattr(layer_mod.current_omni_platform, "is_npu", lambda: True)
+        monkeypatch.setattr(
+            layer_mod.current_omni_platform,
+            "supports_diffusion_dense_flash_attention",
+            lambda: mindiesd_available,
+        )
+
+        result = attention._run_local_attention(
+            torch.zeros(1, dtype=torch.bfloat16),
+            torch.zeros(1, dtype=torch.bfloat16),
+            torch.zeros(1, dtype=torch.bfloat16),
+            AttentionMetadata(),
+        )
+
+        if mindiesd_available:
+            assert result is flash_output
+            dense_flash.assert_called_once()
+            sdpa.assert_not_called()
+        else:
+            assert result is sdpa_output
+            dense_flash.assert_not_called()
+            sdpa.assert_called_once()
+
+    def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch):
+        attention = Attention.__new__(Attention)
+        attention._scheduler_paged_kv = False
+        attention.paged_kv_cache_role = "primary"
+        attention.backend_pref = "FLASH_ATTN"
+        attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
+        attention.attention = SimpleNamespace(
+            forward=Mock(side_effect=ModuleNotFoundError("No module named 'mindiesd'", name="mindiesd"))
+        )
+        attention.sdpa_fallback = SimpleNamespace(forward=Mock(side_effect=AssertionError("must not fall back")))
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+
+        with pytest.raises(ModuleNotFoundError, match="mindiesd"):
+            attention._run_local_attention(
+                torch.zeros(1, dtype=torch.bfloat16),
+                torch.zeros(1, dtype=torch.bfloat16),
+                torch.zeros(1, dtype=torch.bfloat16),
+                AttentionMetadata(),
+            )
+
+    def test_paged_marker_does_not_change_dense_backend(self, monkeypatch):
+        class _FakeAttentionImpl:
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+
+            def forward(self, query, key, value, attn_metadata=None):
+                return query
+
+        class _DenseBackend:
+            supports_paged_kv = False
+
+            @staticmethod
+            def get_name() -> str:
+                return "TORCH_SDPA"
+
+            @staticmethod
+            def get_impl_cls():
+                return _FakeAttentionImpl
+
+        calls = []
+
+        def _fake_get_attn_backend_for_role(*args, **kwargs):
+            calls.append((args, kwargs))
+            return _DenseBackend, None
+
+        monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
+        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
+        monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(),
+            diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.DENSE_LEGACY,
+            parallel_config=SimpleNamespace(ring_degree=1),
+            diffusion_kv_cache_dtype=None,
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config):
+            attention = Attention(
+                num_heads=4,
+                head_size=64,
+                causal=False,
+                softmax_scale=1.0,
+                paged_kv_cache_role="primary",
+            )
+
+        assert len(calls) == 1
+        assert attention.attn_backend is _DenseBackend
+        assert attention.backend_pref is None
+
     def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch):
         class _FakeAttentionImpl:
             def __init__(self, **kwargs):

tests/diffusion/attention/test_piecewise_attn.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
 """End-to-end test for ``piecewise_attn`` (CPU).
 
 Verify that running attention in segments (causal outside full-attn spans,
@@ -98,9 +98,12 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
     query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
     key = query + 100
     value = query + 200
-    calls = []
+    calls: list[object] = []
 
-    def run_segment(segment_query, segment_key, segment_value, segment_batch):
+    output_buffer = torch.empty_like(query)
+
+    def run_segment(segment_query, segment_key, segment_value, segment_batch, segment_output_buffer):
+        assert segment_output_buffer is None
         calls.append(
             (
                 [segment.mode for segment in segment_batch.row_segments],
@@ -119,9 +122,13 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
         plan,
         plan.segments,
         run_segment,
+        output_buffer=output_buffer,
     )
 
+    assert output is output_buffer
     torch.testing.assert_close(output, query)
+    assert all(segment.query_range is None for segment in plan.segments)
+    assert not plan.segments_cover_query_contiguously
     assert calls == [
         (["causal", "causal"], [2, 4], [0, 1, 6, 7]),
         (["full", "full"], [5, 7], [2, 3, 4, 8, 9, 10]),
@@ -129,6 +136,161 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
     ]
 
 
+def test_paged_piecewise_runner_preserves_homogeneous_batch_without_scatter(monkeypatch: pytest.MonkeyPatch):
+    plan = build_paged_piecewise_plan(
+        [[(2, 5)], [(2, 5)]],
+        query_offsets=[0, 0],
+        query_lens=[6, 6],
+        seq_lens=[6, 6],
+        device=DEVICE,
+    )
+    assert plan.homogeneous_batch_shape == (2, 6)
+    query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
+    calls: list[torch.Tensor] = []
+
+    def fail_index_copy(*_args, **_kwargs):
+        raise AssertionError("homogeneous piecewise path must not use index_copy_")
+
+    monkeypatch.setattr(torch.Tensor, "index_copy_", fail_index_copy)
+
+    def run_segment(segment_query, segment_key, segment_value, _metadata, segment_output_buffer):
+        assert segment_key is None
+        assert segment_value is None
+        assert segment_output_buffer is None
+        calls.append(segment_query.clone())
+        return segment_query + 1
+
+    output_buffer = torch.empty_like(query)
+    output = run_paged_piecewise_plan(
+        query,
+        None,
+        None,
+        plan,
+        plan.segments,
+        run_segment,
+        output_buffer=output_buffer,
+        use_homogeneous_batch=True,
+    )
+
+    assert output is output_buffer
+    torch.testing.assert_close(output, query + 1)
+    assert [call.shape[0] for call in calls] == [4, 6, 2]
+    torch.testing.assert_close(calls[0], query[[0, 1, 6, 7]])
+    torch.testing.assert_close(calls[1], query[[2, 3, 4, 8, 9, 10]])
+    torch.testing.assert_close(calls[2], query[[5, 11]])
+
+
+def test_paged_piecewise_runner_keeps_indexed_fallback_for_heterogeneous_layout():
+    plan = build_paged_piecewise_plan(
+        [[(2, 5)], [(1, 4)]],
+        query_offsets=[0, 0],
+        query_lens=[6, 6],
+        seq_lens=[6, 6],
+        device=DEVICE,
+    )
+    assert plan.homogeneous_batch_shape is None
+    query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
+    output = run_paged_piecewise_plan(
+        query,
+        None,
+        None,
+        plan,
+        plan.segments,
+        lambda segment_query, *_args: segment_query,
+        use_homogeneous_batch=True,
+    )
+    torch.testing.assert_close(output, query)
+
+
+def test_paged_piecewise_runner_slices_contiguous_segments():
+    plan = build_paged_piecewise_plan(
+        [[(2, 5)]],
+        query_offsets=[0],
+        query_lens=[6],
+        seq_lens=[6],
+        device=DEVICE,
+    )
+    query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
+    key = query + 100
+    value = query + 200
+    expected_ranges = [(0, 2), (2, 5), (5, 6)]
+    calls: list[tuple[int, int]] = []
+
+    def run_segment(segment_query, segment_key, segment_value, _segment_batch, segment_output_buffer):
+        assert segment_output_buffer is None
+        expected_range = expected_ranges[len(calls)]
+        start, end = expected_range
+        calls.append(expected_range)
+        assert segment_query.untyped_storage().data_ptr() == query.untyped_storage().data_ptr()
+        assert segment_key.untyped_storage().data_ptr() == key.untyped_storage().data_ptr()
+        assert segment_value.untyped_storage().data_ptr() == value.untyped_storage().data_ptr()
+        torch.testing.assert_close(segment_query, query[start:end])
+        return segment_query
+
+    output = run_paged_piecewise_plan(query, key, value, plan, plan.segments, run_segment)
+
+    assert calls == expected_ranges
+    assert [segment.query_range for segment in plan.segments] == expected_ranges
+    assert plan.segments_cover_query_contiguously
+    torch.testing.assert_close(output, query)
+
+
+def test_paged_piecewise_runner_writes_contiguous_segments_into_output_buffer():
+    plan = build_paged_piecewise_plan(
+        [[(2, 5)]],
+        query_offsets=[0],
+        query_lens=[6],
+        seq_lens=[6],
+        device=DEVICE,
+    )
+    query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
+    output_buffer = torch.empty(6, 2, 3)
+    targets = []
+
+    def run_segment(segment_query, _key, _value, _metadata, segment_output_buffer):
+        assert segment_output_buffer is not None
+        assert segment_output_buffer.untyped_storage().data_ptr() == output_buffer.untyped_storage().data_ptr()
+        targets.append(segment_output_buffer)
+        return segment_output_buffer.copy_(segment_query[..., :3])
+
+    output = run_paged_piecewise_plan(
+        query,
+        query,
+        query,
+        plan,
+        plan.segments,
+        run_segment,
+        output_buffer=output_buffer,
+    )
+
+    assert output is output_buffer
+    assert [target.shape[0] for target in targets] == [2, 3, 1]
+    torch.testing.assert_close(output, query[..., :3])
+
+
+def test_paged_piecewise_runner_returns_single_full_segment_directly():
+    plan = build_paged_piecewise_plan(
+        [[]],
+        query_offsets=[0],
+        query_lens=[6],
+        seq_lens=[6],
+        device=DEVICE,
+    )
+    query = torch.randn(6, 2, 4)
+    runner_output = query + 1
+
+    output = run_paged_piecewise_plan(
+        query,
+        query,
+        query,
+        plan,
+        plan.segments,
+        lambda *_args: runner_output,
+    )
+
+    assert output is runner_output
+
+
 def test_paged_piecewise_runner_uses_runner_output_shape():
     plan = build_paged_piecewise_plan(
         [[(2, 5)]],