[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)]],

tests/diffusion/diffusion_kv/test_block_tables.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
 
 from __future__ import annotations
 
@@ -485,8 +485,16 @@ def test_valid_sequence_and_context_install_into_native_rows(monkeypatch: pytest
     assert runner.get_diffusion_kv_row("req-0", 0, "text") == 1
     sequence_binding = runner._resolve_paged_attention_row("req-0", 0, None)
     context_binding = runner._resolve_paged_attention_row("req-0", None, "text")
-    assert (sequence_binding.row_index, sequence_binding.max_seq_len) == (0, 9)
-    assert (context_binding.row_index, context_binding.max_seq_len) == (1, 5)
+    assert (sequence_binding.row_index, sequence_binding.max_seq_len, sequence_binding.block_ids) == (
+        0,
+        9,
+        ((1, 2, 3), (4, 5)),
+    )
+    assert (context_binding.row_index, context_binding.max_seq_len, context_binding.block_ids) == (
+        1,
+        5,
+        ((6, 7), (8,)),
+    )
 
 
 def test_block_table_mutations_invalidate_prepared_attention_batches(

tests/diffusion/diffusion_kv/test_config.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
 
 import pytest
 
@@ -36,7 +36,7 @@ def test_non_mapping_omni_kv_config_is_rejected() -> None:
         OmniDiffusionConfig.from_kwargs(omni_kv_config="paged_scheduler")
 
 
-def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
+def test_paged_scheduler_rejects_dense_connector_kv_receive() -> None:
     with pytest.raises(ValueError, match="does not support imported AR KV"):
         OmniDiffusionConfig.from_kwargs(
             diffusion_kv_mode="paged_scheduler",

tests/diffusion/diffusion_kv/test_initialization.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
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -290,6 +290,7 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None
         dtype=torch.bfloat16,
         non_causal=True,
     )
+    model_limit = vllm_config.model_config.max_model_len
 
     initialization.build_native_kv_cache_configs(
         vllm_config,
@@ -298,4 +299,8 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None
     )
 
     assert vllm_config.model_config.original_max_model_len == -1
-    assert vllm_config.model_config.max_model_len == 256
+    # Sixteen physical pages provide ``16 * block_size`` tokens.  The native
+    # backend chooses the block geometry, so keep this assertion backend
+    # independent instead of assuming the generic 16-token block size.
+    expected_max_model_len = min(model_limit, spec.block_size * 16)
+    assert vllm_config.model_config.max_model_len == expected_max_model_len

tests/diffusion/diffusion_kv/test_paged_attention_adapter.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
 
 from __future__ import annotations
 
@@ -20,12 +20,16 @@ from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionAdapter,
+    DiffusionPagedAttentionMetadata,
     DiffusionPagedAttentionRow,
     DiffusionPagedAttentionRowBinding,
+    DiffusionPagedAttentionRuntime,
 )
 from vllm_omni.diffusion.forward_context import (
+    get_forward_context,
     override_paged_kv_adapter,
     set_forward_context,
+    set_forward_context_denoise_step_idx,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -122,8 +126,8 @@ class _FakeNativeImpl:
     def do_kv_cache_update(
         self,
         _layer,
-        key: torch.Tensor,
-        value: torch.Tensor,
+        key: torch.Tensor | None,
+        value: torch.Tensor | None,
         _kv_cache,
         slot_mapping: torch.Tensor,
     ) -> None:
@@ -155,6 +159,11 @@ def _make_adapter(
     capacity: int = 16,
 ) -> tuple[DiffusionPagedAttentionAdapter, _FakeBlockTables, _FakeLayer, list[tuple]]:
     events: list[tuple] = []
+    monkeypatch.setattr(
+        adapter_module.current_omni_platform,
+        "requires_diffusion_paged_kv_prewrite",
+        lambda: False,
+    )
     block_tables = _FakeBlockTables()
     block_tables.max_num_batched_tokens = capacity
     block_tables.num_blocks.np.fill(math.ceil(capacity / block_tables.block_sizes[0]))
@@ -301,6 +310,47 @@ def test_paged_adapter_accepts_native_gqa_shapes(monkeypatch: pytest.MonkeyPatch
     assert context.value_write.shape == (3, 8, 4)
 
 
+def test_paged_adapter_preserves_strided_kv_for_platform_writer(monkeypatch: pytest.MonkeyPatch) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    batch = adapter.prepare_batch(
+        [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
+    )
+    packed_qkv = torch.randn(1, 5, 3 * 2 * 4)
+    query, key, value = packed_qkv.split(2 * 4, dim=-1)
+    query = query.reshape(1, 5, 2, 4)
+    key = key.reshape(1, 5, 2, 4)
+    value = value.reshape(1, 5, 2, 4)
+    assert not key.is_contiguous()
+    assert not value.is_contiguous()
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context("layer-0", query, key, value)
+
+    assert context.query.is_contiguous()
+    assert not context.key_write.is_contiguous()
+    assert not context.value_write.is_contiguous()
+    assert context.key_write.data_ptr() == key.data_ptr()
+    assert context.value_write.data_ptr() == value.data_ptr()
+
+
+def test_paged_adapter_packs_noncontiguous_kv_head_layout(monkeypatch: pytest.MonkeyPatch) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    batch = adapter.prepare_batch(
+        [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
+    )
+    query = torch.randn(1, 5, 2, 4)
+    key = torch.randn(1, 5, 4, 2).transpose(-1, -2)
+    value = torch.randn(1, 5, 4, 2).transpose(-1, -2)
+    assert key.shape == (1, 5, 2, 4)
+    assert key.stride(-1) != 1
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context("layer-0", query, key, value)
+
+    assert context.key_write.is_contiguous()
+    assert context.value_write.is_contiguous()
+
+
 def test_omni_paged_backend_defers_backend_owned_cache_update(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, layer, _ = _make_adapter(monkeypatch)
     layer.attn_backend.forward_includes_kv_cache_update = True
@@ -350,10 +400,48 @@ def test_omni_piecewise_backend_defers_backend_owned_cache_update(monkeypatch: p
     assert [metadata.query_start_loc_cpu.tolist() for metadata in segment_metadata] == [[0, 1], [0, 2], [0, 1]]
     assert [metadata.positions.tolist() for metadata in segment_metadata] == [[0], [1, 2], [3]]
     assert [metadata.slot_mappings.tolist() for metadata in segment_metadata] == [[[0]], [[1, 2]], [[3]]]
+    for event in events[1:]:
+        build_kwargs = event[1]
+        assert build_kwargs["positions"].untyped_storage().data_ptr() == batch.positions.untyped_storage().data_ptr()
+        assert (
+            build_kwargs["slot_mappings"].untyped_storage().data_ptr()
+            == batch.slot_mappings.untyped_storage().data_ptr()
+        )
+
+
+def test_platform_prewrite_updates_once_then_piecewise_reads_cache(monkeypatch: pytest.MonkeyPatch) -> None:
+    adapter, _, layer, _ = _make_adapter(monkeypatch)
+    layer.attn_backend.forward_includes_kv_cache_update = True
+    monkeypatch.setattr(
+        adapter_module.current_omni_platform,
+        "requires_diffusion_paged_kv_prewrite",
+        lambda: True,
+    )
+    adapter.resolve_row = lambda *_args: DiffusionPagedAttentionRowBinding(
+        row_index=2,
+        max_seq_len=16,
+        block_ids=((3, 7, 9, 11),),
+    )
+    row = DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)
+    batch = adapter.prepare_batch([row])
+    qkv = torch.randn(1, 4, 2, 4)
+    metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv, omni_attn_metadata=metadata)
+        output = _run_omni_paged_backend(context)
+
+    assert torch.equal(output, qkv)
+    assert layer.native_events == ["update", "forward", "forward", "forward"]
+    assert len(layer.updates) == 1
+    assert all(call[1] is None and call[2] is None for call in layer.calls)
 
 
 def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, layer, events = _make_adapter(monkeypatch)
+    # Exercise the Ascend-only homogeneous batch path. CUDA keeps its native
+    # output-buffer contract and is covered by the GPU adapter tests.
+    monkeypatch.setattr(adapter_module.current_omni_platform, "is_npu", lambda: True)
     batch = adapter.prepare_batch(
         [
             DiffusionPagedAttentionRow(
@@ -450,6 +538,48 @@ def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.
     ]
 
 
+def test_piecewise_metadata_snapshots_reused_native_scheduler_buffer(monkeypatch: pytest.MonkeyPatch) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    shared_scheduler_metadata = torch.zeros(1, dtype=torch.int32)
+    build_id = 0
+
+    def build_metadata(**kwargs):
+        nonlocal build_id
+        shared_scheduler_metadata.fill_(build_id)
+        metadata = SimpleNamespace(
+            build_id=build_id,
+            causal=kwargs["causal"],
+            seq_lens=kwargs["seq_lens"].clone(),
+            query_start_loc_cpu=kwargs["query_start_loc_cpu"].clone(),
+            positions=kwargs["positions"].clone(),
+            slot_mappings=kwargs["slot_mappings"].clone(),
+            scheduler_metadata=shared_scheduler_metadata,
+            prefix_scheduler_metadata=None,
+        )
+        build_id += 1
+        return {"layer-0": metadata}
+
+    monkeypatch.setattr(adapter_module.current_omni_platform, "build_diffusion_kv_attn_metadata", build_metadata)
+    batch = adapter.prepare_batch(
+        [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
+    )
+    qkv = torch.randn(1, 4, 2, 4)
+    metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context(
+            "layer-0",
+            qkv,
+            qkv,
+            qkv,
+            omni_attn_metadata=metadata,
+        )
+
+    segment_schedules = [item.scheduler_metadata for item in context.piecewise_native_metadata]
+    assert [item.item() for item in segment_schedules] == [1, 2, 3]
+    assert len({item.data_ptr() for item in segment_schedules}) == 3
+
+
 def test_omni_paged_backend_treats_empty_full_spans_as_causal(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, layer, events = _make_adapter(monkeypatch)
     batch = adapter.prepare_batch(
@@ -851,6 +981,28 @@ def test_forward_rejects_batched_layout_that_does_not_match_rows(monkeypatch: py
         adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
 
 
+def test_forward_packs_and_restores_right_padded_heterogeneous_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    batch = adapter.prepare_batch(
+        [
+            DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1),
+            DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=3, seq_len=3),
+        ]
+    )
+    qkv = torch.arange(2 * 3 * 2 * 4, dtype=torch.float32).reshape(2, 3, 2, 4)
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
+
+    assert context.output_scatter_indices.tolist() == [0, 3, 4, 5]
+    torch.testing.assert_close(context.query, qkv.reshape(-1, 2, 4)[[0, 3, 4, 5]])
+    restored = context.restore_output(context.query)
+    assert restored.shape == qkv.shape
+    torch.testing.assert_close(restored[0, 0], qkv[0, 0])
+    assert torch.count_nonzero(restored[0, 1:]) == 0
+    torch.testing.assert_close(restored[1], qkv[1])
+
+
 def test_forward_rejects_padded_full_kv(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, _, _ = _make_adapter(monkeypatch)
     batch = adapter.prepare_batch(
@@ -878,6 +1030,41 @@ def test_forward_rejects_padded_full_kv(monkeypatch: pytest.MonkeyPatch) -> None
         adapter.prepare_layer_context("layer-0", query, full_kv, full_kv)
 
 
+def test_runner_runtime_switches_from_prefill_to_reusable_denoise_metadata(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    metadata = DiffusionPagedAttentionMetadata(
+        prefill_rows=(DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=6, seq_len=6),),
+        denoise_rows=(
+            DiffusionPagedAttentionRow(
+                request_id="req-0",
+                sequence_id=0,
+                query_len=2,
+                seq_len=5,
+                kv_start_pos=3,
+            ),
+        ),
+    )
+    runtime = DiffusionPagedAttentionRuntime(adapter, metadata)
+
+    with set_forward_context(paged_kv_runtime=runtime), runtime.activate():
+        set_forward_context_denoise_step_idx(0)
+        assert Attention._active_paged_kv_adapter() is adapter
+        assert adapter._active_batch.rows == metadata.prefill_rows
+
+        set_forward_context_denoise_step_idx(1)
+        assert Attention._active_paged_kv_adapter() is adapter
+        denoise_batch = adapter._active_batch
+        assert denoise_batch.rows == metadata.denoise_rows
+
+        set_forward_context_denoise_step_idx(2)
+        assert Attention._active_paged_kv_adapter() is adapter
+        assert adapter._active_batch is denoise_batch
+
+    assert adapter._active_batch is None
+
+
 def test_causal_batch_rejects_non_suffix_write(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, _, _ = _make_adapter(monkeypatch, non_causal=False)
     row = DiffusionPagedAttentionRow(
@@ -992,7 +1179,8 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     selected_backends = []
-    impl_cls = Mock(return_value=SimpleNamespace(forward=Mock()))
+    specialized_backends = []
+    impl_cls = Mock(return_value=SimpleNamespace(forward=Mock(), do_kv_cache_update=Mock()))
     native_backend = SimpleNamespace(
         get_name=lambda: "ASCEND",
         indexes_kv_by_block_stride=lambda: True,
@@ -1003,16 +1191,28 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
     original_backend_per_kind = {"full": object()}
     config = SimpleNamespace(
         attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind),
+        parallel_config=SimpleNamespace(prefill_context_parallel_size=2),
         model_config=SimpleNamespace(dtype=torch.float16),
         cache_config=SimpleNamespace(cache_dtype="auto"),
     )
 
     def select_backend(**_kwargs):
-        selected_backends.append((config.attention_config.backend, config.attention_config.backend_per_kind))
+        selected_backends.append(
+            (
+                config.attention_config.backend,
+                config.attention_config.backend_per_kind,
+                config.parallel_config.prefill_context_parallel_size,
+            )
+        )
         return native_backend
 
     monkeypatch.setattr(adapter_module, "get_attn_backend", select_backend)
     monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
+    monkeypatch.setattr(
+        adapter_module.current_omni_platform,
+        "get_diffusion_paged_kv_attn_backend",
+        lambda backend, *, ulysses_degree: specialized_backends.append((backend, ulysses_degree)) or backend,
+    )
     layer = SimpleNamespace(num_heads=8, softmax_scale=0.125)
     spec = FullAttentionSpec(
         block_size=16,
@@ -1031,9 +1231,11 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         ulysses_degree=2,
     )
 
-    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {})]
+    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 2)]
+    assert specialized_backends == [(native_backend, 2)]
     assert config.attention_config.backend is None
     assert config.attention_config.backend_per_kind is original_backend_per_kind
+    assert config.parallel_config.prefill_context_parallel_size == 2
     assert native_layer.num_heads == 4
     assert native_layer.num_kv_heads == 2
     assert native_layer.spec.num_kv_heads == 2
@@ -1094,6 +1296,51 @@ def test_omni_attention_wraps_paged_kernel_with_sp_hooks() -> None:
     assert torch.equal(output, torch.full_like(original, 6))
 
 
+def test_omni_attention_strips_paged_ulysses_padding_around_kernel() -> None:
+    events: list[tuple[str, tuple[int, ...]]] = []
+
+    class Strategy:
+        name = "ulysses"
+
+        def pre_attention(self, query, key, value, metadata):
+            return query, key, value, metadata, SimpleNamespace(joint_len=0, joint_strategy="front")
+
+        def post_attention(self, output, _ctx):
+            events.append(("post", tuple(output.shape)))
+            return output
+
+    class Adapter:
+        def prepare_layer_context(self, _layer_name, query, key, value, *, omni_attn_metadata):
+            assert omni_attn_metadata is None
+            events.append(("kernel_input", tuple(query.shape)))
+            assert query[:, -1].tolist() == [[[3.0]]]
+            assert torch.equal(query, key)
+            assert torch.equal(query, value)
+            return SimpleNamespace(query=query)
+
+    class Backend:
+        def forward_paged(self, context):
+            return context.query
+
+    layer = Attention.__new__(Attention)
+    nn.Module.__init__(layer)
+    layer.prefix = "layer-0"
+    layer.paged_kv_cache_role = "primary"
+    layer.attn_backend = SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN")
+    layer.attention = Backend()
+    layer.use_ring = False
+    layer._no_parallel_strategy = object()
+    layer._get_active_parallel_strategy = lambda: Strategy()
+    tensor = torch.arange(5, dtype=torch.float32).reshape(1, 5, 1, 1)
+
+    with set_forward_context(), override_paged_kv_adapter(Adapter()):
+        get_forward_context().sp_padding_size = 1
+        output = layer._forward_impl(tensor, tensor, tensor)
+
+    assert events == [("kernel_input", (1, 4, 1, 1)), ("post", (1, 5, 1, 1))]
+    assert output.flatten().tolist() == [0.0, 1.0, 2.0, 3.0, 0.0]
+
+
 def test_omni_attention_rejects_paged_request_for_backend_without_paged_support() -> None:
     class Adapter:
         def prepare_layer_context(self, *args, **kwargs):
@@ -1145,3 +1392,15 @@ def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None:
 
     assert events == ["pre", "dense", "post"]
     assert output is qkv
+
+
+def test_omni_attention_rejects_scheduler_paged_forward_without_active_adapter() -> None:
+    layer = Attention.__new__(Attention)
+    nn.Module.__init__(layer)
+    layer._scheduler_paged_kv = True
+    layer.paged_kv_cache_role = "primary"
+    layer._get_active_parallel_strategy = lambda: object()
+    qkv = torch.zeros(1, 2, 2, 4)
+
+    with set_forward_context(), pytest.raises(RuntimeError, match="without an active Worker adapter"):
+        layer._forward_impl(qkv, qkv, qkv)

tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py

@@ -154,8 +154,17 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
 
     torch.manual_seed(1)
     query = torch.randn(19, _NUM_HEADS, _HEAD_SIZE, dtype=torch.float16, device=device)
-    key = torch.randn_like(query)
-    value = torch.randn_like(query)
+    packed_qkv = torch.randn(
+        19,
+        3 * _NUM_HEADS * _HEAD_SIZE,
+        dtype=torch.float16,
+        device=device,
+    )
+    _, key_flat, value_flat = packed_qkv.split(_NUM_HEADS * _HEAD_SIZE, dim=-1)
+    key = key_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
+    value = value_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
+    assert not key.is_contiguous()
+    assert not value.is_contiguous()
     omni_backend = FlashAttentionImpl(
         num_heads=_NUM_HEADS,
         head_size=_HEAD_SIZE,

tests/diffusion/diffusion_kv/test_spec_discovery.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
 
 from contextlib import nullcontext
 from types import SimpleNamespace
@@ -42,9 +42,10 @@ class _RunnerKVBackend:
         self.kv_cache_config = config
 
 
-def _attention(*, enabled: bool) -> Attention:
+def _attention(*, enabled: bool, prefix: str = "image_attention") -> Attention:
     attention = Attention.__new__(Attention)
     nn.Module.__init__(attention)
+    attention.prefix = prefix
     attention.paged_kv_cache_role = "primary" if enabled else None
     attention.paged_kv_cache_dtype = torch.bfloat16
     attention.num_kv_heads = 2
@@ -82,6 +83,23 @@ def test_runner_discovers_native_spec_from_loaded_attention() -> None:
     assert spec.non_causal is True
 
 
+def test_runner_uses_attention_prefix_as_canonical_layer_name() -> None:
+    runner = _runner(_attention(enabled=True, prefix="layers.0.self_attn.image_attn.attn"))
+
+    specs = runner.get_kv_cache_spec()
+
+    assert set(specs) == {"layers.0.self_attn.image_attn.attn"}
+
+
+def test_runner_rejects_duplicate_attention_prefixes() -> None:
+    runner = _runner(_attention(enabled=True, prefix="shared"))
+    second_attention = _attention(enabled=True, prefix="shared")
+    runner.pipeline.second_image_attention = second_attention
+
+    with pytest.raises(RuntimeError, match="Duplicate canonical paged Diffusion Attention prefix 'shared'"):
+        runner.get_kv_cache_spec()
+
+
 def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None:
     runner = _runner(_attention(enabled=False))
 

tests/diffusion/diffusion_kv/test_worker_contract.py

@@ -42,7 +42,10 @@ def make_metadata(request_id: str = "req-0") -> DiffusionKVMetadata:
 
 def make_runner(mode: DiffusionKVCacheMode) -> DiffusionModelRunner:
     runner = object.__new__(DiffusionModelRunner)
-    runner.od_config = SimpleNamespace(diffusion_kv_mode=mode)
+    runner.od_config = SimpleNamespace(
+        diffusion_kv_mode=mode,
+        parallel_config=SimpleNamespace(cfg_parallel_size=1),
+    )
     return runner
 
 
@@ -85,6 +88,51 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None:
     assert new_req.diffusion_kv_metadata is metadata
 
 
+def test_runner_builds_prefill_and_denoise_rows_from_scheduler_metadata() -> None:
+    runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
+    metadata = DiffusionKVMetadata(
+        request_id="req-0",
+        allocation_generation=1,
+        sequences=(
+            DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)),
+            DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)),
+        ),
+    )
+
+    attn_metadata = runner._build_paged_attention_metadata([metadata])
+
+    assert [
+        (row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len)
+        for row in attn_metadata.prefill_rows
+    ] == [("req-0", 0, 0, 8, 8), ("req-0", 1, 0, 9, 9)]
+    assert [
+        (row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len)
+        for row in attn_metadata.denoise_rows
+    ] == [("req-0", 0, 4, 2, 6), ("req-0", 1, 5, 2, 7)]
+
+
+def test_runner_selects_only_local_cfg_parallel_row(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
+    runner.od_config.parallel_config.cfg_parallel_size = 2
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.worker.diffusion_model_runner.get_classifier_free_guidance_rank",
+        lambda: 1,
+    )
+    metadata = DiffusionKVMetadata(
+        request_id="req-0",
+        allocation_generation=1,
+        sequences=(
+            DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)),
+            DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)),
+        ),
+    )
+
+    attn_metadata = runner._build_paged_attention_metadata([metadata])
+
+    assert [row.sequence_id for row in attn_metadata.prefill_rows] == [1]
+    assert [row.sequence_id for row in attn_metadata.denoise_rows] == [1]
+
+
 def test_paged_request_without_metadata_fails_before_request_forward() -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
     runner._execute_request_list = Mock()

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.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
 
 from types import SimpleNamespace
 
@@ -270,6 +270,17 @@ def test_grouped_denoise_allows_sdpa_attention_backend():
     pipeline._ensure_grouped_attention_backend_supported(2)
 
 
+def test_scheduler_paged_step_execution_is_rejected():
+    pipeline = _pipeline()
+    pipeline.od_config.diffusion_kv_mode = hy3_module.DiffusionKVCacheMode.PAGED_SCHEDULER
+    state = _state("paged-step", 0)
+
+    with pytest.raises(ValueError, match="request-level execution only"):
+        pipeline.prepare_encode(state)
+    with pytest.raises(ValueError, match="request-level execution only"):
+        pipeline.denoise_step(InputBatch.make_batch([state]))
+
+
 def test_step_scheduler_preserves_latent_dtype_for_mixed_progress_batches():
     pipeline = _pipeline()
     pipeline._pipeline = SimpleNamespace(prepare_extra_func_kwargs=lambda step, kwargs: {})

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.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 ImageKVCacheManager.
 
 Covers: cache → reuse flow, AR KV injection, CFG (sequential & parallel), SP, cross-request isolation.
@@ -37,12 +37,21 @@ class MockAttention(nn.Module):
         super().__init__()
         self.paged_kv_active = False
         self.calls = []
+        self.paged_calls = []
+        self.paged_metadata = []
 
     def is_paged_kv_active(self):
         return self.paged_kv_active
 
     def forward(self, query, key, value, attn_metadata=None, **kwargs):
         self.calls.append((query, key, value))
+        from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+        if is_forward_context_available() and get_forward_context().paged_kv_adapter is not None:
+            self.paged_calls.append((query, key, value))
+            self.paged_metadata.append(attn_metadata)
+        if attn_metadata is not None and attn_metadata.joint_query is not None:
+            return torch.cat((attn_metadata.joint_query, query), dim=1)
         return query
 
 
@@ -105,6 +114,7 @@ def _call_mgr(
     shard_image_size=None,
     gen_timestep_scatter_index=None,
     position_ids=None,
+    full_attn_spans=None,
 ):
     query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM)
     attn_mask = torch.zeros(bs, 1, seq_len, seq_len)
@@ -121,6 +131,7 @@ def _call_mgr(
         shard_image_size=shard_image_size,
         gen_timestep_scatter_index=gen_timestep_scatter_index,
         position_ids=position_ids,
+        full_attn_spans=full_attn_spans,
     )
 
 
@@ -132,6 +143,83 @@ def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
     assert mgr.state_dict() == {}
 
 
+def test_scheduler_paged_kv_runs_piecewise_for_first_and_later_steps() -> None:
+    from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
+
+    mgr = _make_cache_mgr()
+    mgr.attn.paged_kv_active = True
+    bs = 2
+    prefix_lens = torch.tensor([[2], [4]], dtype=torch.long)
+    first_q_len = 12
+    first_k, first_v = _make_known_kv(bs * first_q_len, base=1.0)
+    full_attn_spans = [[(2, 12)], [(2, 12)]]
+
+    with set_forward_context(), override_paged_kv_adapter(object()):
+        _call_mgr(
+            mgr,
+            bs,
+            q_len=first_q_len,
+            seq_len=first_q_len,
+            key_flat=first_k,
+            value_flat=first_v,
+            first_step=True,
+            gen_timestep_scatter_index=prefix_lens,
+            full_attn_spans=full_attn_spans,
+        )
+
+        assert mgr.image_kv_cache_map is None
+        assert mgr.image_kv_cache_lens is None
+        assert len(mgr.attn.paged_calls) == 1
+        assert mgr.attn.paged_calls[0][0].shape == (bs, first_q_len, NUM_HEADS, HEAD_DIM)
+        assert mgr.attn.paged_calls[0][1].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM)
+        assert mgr.attn.paged_calls[0][2].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM)
+        assert mgr.attn.paged_metadata[0].full_attn_spans == full_attn_spans
+
+        current_q_len = IMAGE_TOKEN_LEN
+        current_k, current_v = _make_known_kv(bs * current_q_len, base=50.0)
+        _call_mgr(
+            mgr,
+            bs,
+            q_len=current_q_len,
+            seq_len=prefix_lens[-1].item() + current_q_len,
+            key_flat=current_k,
+            value_flat=current_v,
+            position_ids=torch.arange(4, 4 + current_q_len).repeat(bs, 1),
+            full_attn_spans=full_attn_spans,
+        )
+
+    assert len(mgr.attn.paged_calls) == 2
+    assert mgr.attn.paged_calls[1][0].shape == (bs, current_q_len, NUM_HEADS, HEAD_DIM)
+    assert mgr.attn.paged_calls[1][1].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM)
+    assert mgr.attn.paged_calls[1][2].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM)
+    assert mgr.attn.paged_metadata[1].full_attn_spans == full_attn_spans
+
+
+def test_cache_manager_rejects_dense_imported_prefix_for_paged_attention() -> None:
+    from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
+
+    mgr = _make_cache_mgr()
+    mgr.attn.paged_kv_active = True
+    prefix_len = 4
+    prefix_key, prefix_value = _make_known_kv(prefix_len, base=100.0)
+    mgr._injected_ar_kv = [(prefix_key, prefix_value)]
+    current_key, current_value = _make_known_kv(3, base=1.0)
+
+    with set_forward_context(), override_paged_kv_adapter(object()):
+        with pytest.raises(NotImplementedError, match="does not support imported AR KV"):
+            _call_mgr(
+                mgr,
+                bs=1,
+                q_len=3,
+                seq_len=7,
+                key_flat=current_key,
+                value_flat=current_value,
+                first_step=True,
+                gen_timestep_scatter_index=torch.tensor([[0]]),
+                full_attn_spans=[[(4, 7)]],
+            )
+
+
 @pytest.mark.parametrize(
     ("paged_kv_active", "expected_kv_heads"),
     [(False, NUM_HEADS), (True, NUM_KV_HEADS)],
@@ -155,6 +243,7 @@ def test_gqa_kv_stays_compressed_for_paged_attention(
         first_step=True,
         num_image_tokens=0,
         gen_timestep_scatter_index=_gen_timestep_index(1, 0),
+        full_attn_spans=[[(0, q_len)]] if paged_kv_active else None,
     )
 
     _, key_input, value_input = mgr.attn.calls[-1]

tests/diffusion/test_diffusion_model_runner.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
 
 from contextlib import contextmanager
 from types import SimpleNamespace
@@ -533,7 +533,14 @@ def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatc
         return original_execute(*args, **kwargs)
 
     runner._execute_request_list = execute
-    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    forward_context_calls = []
+
+    @contextmanager
+    def record_forward_context(*args, **kwargs):
+        forward_context_calls.append((args, kwargs))
+        yield
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context)
     reset_peak_memory_stats = Mock()
     monkeypatch.setattr(
         model_runner_module.current_omni_platform,
@@ -549,6 +556,8 @@ def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatc
 
     assert runner.pipeline.forward_calls == 1
     assert record_names == ["pipeline_memory_profile"]
+    assert len(forward_context_calls) == 1
+    assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True
     runner._validate_diffusion_kv_metadata.assert_not_called()
     reset_peak_memory_stats.assert_not_called()
     synchronize.assert_called_once_with()
@@ -575,7 +584,14 @@ def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypa
 
     runner.pipeline.denoise_step = denoise_step
     runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
-    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    forward_context_calls = []
+
+    @contextmanager
+    def record_forward_context(*args, **kwargs):
+        forward_context_calls.append((args, kwargs))
+        yield
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context)
     reset_peak_memory_stats = Mock()
     monkeypatch.setattr(
         model_runner_module.current_omni_platform,
@@ -592,6 +608,8 @@ def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypa
     assert observed_batch_rows == [(2, 2)]
     assert runner.state_cache == {}
     assert runner.input_batch is None
+    assert len(forward_context_calls) == 1
+    assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True
     runner._validate_diffusion_kv_metadata.assert_not_called()
     reset_peak_memory_stats.assert_not_called()
     synchronize.assert_called_once_with()

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

vllm_omni/diffusion/attention/backends/flash_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
 
 import os
 from functools import partial
@@ -7,11 +7,7 @@ from functools import partial
 import torch
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.attention.backends.abstract import (
-    AttentionBackend,
-    AttentionImpl,
-    AttentionMetadata,
-)
+from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata
 from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
 from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
@@ -285,7 +281,9 @@ class FlashAttentionImpl(AttentionImpl):
         rank-local native cache context.  The adapter no longer performs the
         attention call itself; this method is the backend-owned execution
         boundary.  Its native layer wrapper keeps vLLM version-specific cache
-        and kernel details out of Omni's common ``Attention`` layer.
+        and kernel details out of Omni's common ``Attention`` layer. CUDA uses
+        the native vLLM FlashAttention writer/kernel contract. Ascend writes
+        the complete K/V span once before its piecewise FIA calls.
         """
 
         layer = getattr(paged_kv_context, "layer", None)
@@ -297,8 +295,18 @@ class FlashAttentionImpl(AttentionImpl):
         native_impl = getattr(layer, "impl", None)
         if native_impl is None:
             raise RuntimeError(f"Native attention implementation is not bound for diffusion layer {layer.layer_name!r}")
-        if not layer.attn_backend.forward_includes_kv_cache_update:
-            native_impl.do_kv_cache_update(
+        prewrite_kv = current_omni_platform.requires_diffusion_paged_kv_prewrite()
+        read_kv_from_cache = prewrite_kv and layer.attn_backend.forward_includes_kv_cache_update
+        # The GPU/default path preserves native cache-update ownership. Ascend
+        # prewrites through vLLM-Ascend's normal-layout writer so piecewise FIA
+        # segments do not repeatedly scatter the same layer K/V.
+        if not layer.attn_backend.forward_includes_kv_cache_update or prewrite_kv:
+            cache_update = getattr(native_impl, "do_kv_cache_update", None)
+            if not callable(cache_update):
+                raise RuntimeError(
+                    f"Native attention implementation for {layer.layer_name!r} cannot update the KV cache"
+                )
+            cache_update(
                 layer,
                 paged_kv_context.key_write,
                 paged_kv_context.value_write,
@@ -308,33 +316,52 @@ class FlashAttentionImpl(AttentionImpl):
 
         def run_native_attention(
             query: torch.Tensor,
-            key: torch.Tensor,
-            value: torch.Tensor,
+            key: torch.Tensor | None,
+            value: torch.Tensor | None,
             native_metadata,
+            output: torch.Tensor | None,
         ) -> torch.Tensor:
-            output = torch.empty(
-                (query.shape[0], layer.num_heads, layer.head_size_v),
-                dtype=query.dtype,
-                device=query.device,
-            )
+            if output is None:
+                output = torch.empty(
+                    (query.shape[0], layer.num_heads, layer.head_size_v),
+                    dtype=query.dtype,
+                    device=query.device,
+                )
             return native_impl.forward(
                 layer,
                 query,
-                key,
-                value,
+                None if read_kv_from_cache else key,
+                None if read_kv_from_cache else value,
                 kv_cache,
                 native_metadata,
                 output,
             )
 
         if paged_kv_context.piecewise_plan is not None:
+            # Ascend FIA can execute identical CFG rows as one batched call
+            # per piece.  Keep that batch layout through the piecewise runner
+            # so it can concatenate row-major results instead of emitting an
+            # indexed ScatterUpdate for every segment.  The CUDA path keeps
+            # its existing output-buffer contract (including graph capture).
+            use_homogeneous_batch = (
+                current_omni_platform.is_npu() and paged_kv_context.piecewise_plan.homogeneous_batch_shape is not None
+            )
+            output = None
+            if not use_homogeneous_batch:
+                output = torch.empty(
+                    (paged_kv_context.query.shape[0], layer.num_heads, layer.head_size_v),
+                    dtype=paged_kv_context.query.dtype,
+                    device=paged_kv_context.query.device,
+                )
             output = run_paged_piecewise_plan(
                 paged_kv_context.query,
-                paged_kv_context.key_write,
-                paged_kv_context.value_write,
+                None if read_kv_from_cache else paged_kv_context.key_write,
+                None if read_kv_from_cache else paged_kv_context.value_write,
                 paged_kv_context.piecewise_plan,
                 paged_kv_context.piecewise_native_metadata,
                 run_native_attention,
+                output_buffer=output,
+                use_homogeneous_batch=use_homogeneous_batch,
             )
         else:
             output = run_native_attention(
@@ -342,6 +369,7 @@ class FlashAttentionImpl(AttentionImpl):
                 paged_kv_context.key_write,
                 paged_kv_context.value_write,
                 paged_kv_context.native_metadata,
+                None,
             )
         return paged_kv_context.restore_output(output)
 
@@ -516,15 +544,8 @@ class FlashAttentionImpl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
-        try:
-            from mindiesd import attention_forward
-        except ImportError:
-            raise ImportError(
-                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
-                "Please install MindIE-SD to enable NPU attention support. "
-                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
-                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
-            )
+        from mindiesd import attention_forward
+
         # Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the
         # model marks extra["npu_attn_varlen"] and carries packed metadata, so
         # the padding document is excluded without reading or materializing
@@ -633,15 +654,8 @@ class FlashAttentionImpl(AttentionImpl):
             return None
         seq_q, seq_k = resolved
 
-        try:
-            from mindiesd import attention_forward_varlen
-        except ImportError:
-            raise ImportError(
-                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
-                "Please install MindIE-SD to enable NPU attention support. "
-                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
-                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
-            )
+        from mindiesd import attention_forward_varlen
+
         q = query.squeeze(0)  # [T, N, D] == TND
         k = key.squeeze(0)
         v = value.squeeze(0)
@@ -689,15 +703,8 @@ class FlashAttentionImpl(AttentionImpl):
         _, seq_k = resolved
         used_k = seq_k[0]  # real document length (first cumulative end)
 
-        try:
-            from mindiesd import attention_forward
-        except ImportError:
-            raise ImportError(
-                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
-                "Please install MindIE-SD to enable NPU attention support. "
-                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
-                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
-            )
+        from mindiesd import attention_forward
+
         # mindiesd always takes BSND input regardless of `layout`; the arg
         # selects the op-internal layout and laser only supports BNSD, so do
         # NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here.

vllm_omni/diffusion/attention/backends/utils/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
 """
 Piecewise attention for mixed causal / full (bidirectional) masks.
 
@@ -92,6 +92,7 @@ def piecewise_attn(
 ):
     _check_homogeneous(full_attn_spans)
     spans = full_attn_spans[0]
+    ranges: tuple[tuple[int, int, int], ...]
     if query_ranges is None:
         query_len = query.shape[1]
         ranges = ((0, query_len, key.shape[1] - query_len),)
@@ -133,6 +134,11 @@ class PagedPiecewiseSegment:
 
     row_segments: tuple[Segment, ...]
     query_indices: torch.Tensor
+    query_range: tuple[int, int] | None
+    # Local [start, end) range shared by every row when the batch can retain
+    # its original [B, S] layout.  The packed query_indices remain available
+    # for metadata construction and the heterogeneous fallback.
+    local_query_range: tuple[int, int] | None = None
 
 
 @dataclass(frozen=True, slots=True)
@@ -142,10 +148,16 @@ class PagedPiecewisePlan:
     spans: tuple[tuple[tuple[int, int], ...], ...]
     segments: tuple[PagedPiecewiseSegment, ...]
     num_query_tokens: int
+    segments_cover_query_contiguously: bool
+    # (batch_size, query_len) when all rows have the same local segment
+    # layout.  Native FIA still receives flattened row-major tokens, but the
+    # runner can restore outputs with one batch concat instead of one indexed
+    # scatter per segment.
+    homogeneous_batch_shape: tuple[int, int] | None = None
 
 
 PagedPiecewiseRunner = Callable[
-    [torch.Tensor, torch.Tensor, torch.Tensor, object],
+    [torch.Tensor, torch.Tensor | None, torch.Tensor | None, object, torch.Tensor | None],
     torch.Tensor,
 ]
 
@@ -164,7 +176,9 @@ def build_paged_piecewise_plan(
     if row_count == 0 or not (row_count == len(query_offsets) == len(query_lens) == len(seq_lens)):
         raise ValueError("Paged piecewise inputs must have one entry per row")
 
-    spans = tuple(tuple(tuple(span) for span in row_spans) for row_spans in full_attn_spans)
+    spans: tuple[tuple[tuple[int, int], ...], ...] = tuple(
+        tuple((span_start, span_end) for span_start, span_end in row_spans) for row_spans in full_attn_spans
+    )
 
     packed_offsets = [0]
     for row_spans, query_offset, query_len, seq_len in zip(spans, query_offsets, query_lens, seq_lens, strict=True):
@@ -188,46 +202,206 @@ def build_paged_piecewise_plan(
     for row_segments in zip(*segments_by_row, strict=True):
         if len({segment.mode for segment in row_segments}) != 1:
             raise ValueError("Paged piecewise rows must use the same attention mode per segment")
-        query_indices = []
+        query_indices: list[int] = []
+        local_ranges: list[tuple[int, int]] = []
         for row_index, (segment, query_offset) in enumerate(zip(row_segments, query_offsets, strict=True)):
             local_start = segment.q_start - query_offset
             local_end = segment.q_end - query_offset
+            local_ranges.append((local_start, local_end))
             query_indices.extend(range(packed_offsets[row_index] + local_start, packed_offsets[row_index] + local_end))
+        query_range = None
+        if query_indices and query_indices == list(range(query_indices[0], query_indices[-1] + 1)):
+            query_range = (query_indices[0], query_indices[-1] + 1)
+        local_query_range = local_ranges[0] if len(set(local_ranges)) == 1 else None
         packed_segments.append(
             PagedPiecewiseSegment(
                 row_segments=tuple(row_segments),
                 query_indices=torch.tensor(query_indices, dtype=torch.long, device=device),
+                query_range=query_range,
+                local_query_range=local_query_range,
             )
         )
 
-    return PagedPiecewisePlan(spans, tuple(packed_segments), packed_offsets[-1])
+    covered = 0
+    for segment in packed_segments:
+        if segment.query_range is None or segment.query_range[0] != covered:
+            break
+        covered = segment.query_range[1]
+    segments_cover_query_contiguously = covered == packed_offsets[-1]
+    homogeneous_batch_shape: tuple[int, int] | None = None
+    if (
+        row_count > 1
+        and len(set(query_lens)) == 1
+        and all(segment.local_query_range is not None for segment in packed_segments)
+    ):
+        # A common local range is sufficient even when rows have different
+        # global offsets or KV lengths; those values remain row-specific in
+        # native metadata.
+        homogeneous_batch_shape = (row_count, int(query_lens[0]))
+    return PagedPiecewisePlan(
+        spans=spans,
+        segments=tuple(packed_segments),
+        num_query_tokens=packed_offsets[-1],
+        segments_cover_query_contiguously=segments_cover_query_contiguously,
+        homogeneous_batch_shape=homogeneous_batch_shape,
+    )
+
+
+def _run_homogeneous_paged_piecewise_plan(
+    query: torch.Tensor,
+    key: torch.Tensor | None,
+    value: torch.Tensor | None,
+    plan: PagedPiecewisePlan,
+    segment_metadata: Sequence[object],
+    segment_runner: PagedPiecewiseRunner,
+    output_buffer: torch.Tensor | None,
+) -> torch.Tensor:
+    """Run a homogeneous batch while retaining the old [B, S] layout.
+
+    The native FIA contract is flattened TND, so each segment is flattened
+    only for the kernel call.  Results are reshaped back to [B, L] and
+    concatenated along the sequence dimension, which avoids the per-segment
+    ``index_copy_`` used by heterogeneous packed rows.
+    """
+
+    batch_shape = plan.homogeneous_batch_shape
+    assert batch_shape is not None
+    batch_size, query_len = batch_shape
+    if query.shape[0] != batch_size * query_len:
+        raise ValueError(
+            f"Homogeneous paged piecewise plan expects {batch_size * query_len} query tokens, got {query.shape[0]}"
+        )
+    query_rows = query.reshape(batch_size, query_len, *query.shape[1:])
+    key_rows = None if key is None else key.reshape(batch_size, query_len, *key.shape[1:])
+    value_rows = None if value is None else value.reshape(batch_size, query_len, *value.shape[1:])
+
+    segment_outputs: list[torch.Tensor] = []
+    for segment, metadata in zip(plan.segments, segment_metadata, strict=True):
+        local_range = segment.local_query_range
+        if local_range is None:
+            raise ValueError("Homogeneous paged piecewise execution requires a shared local segment range")
+        start, end = local_range
+        segment_query = query_rows[:, start:end].contiguous().reshape(-1, *query.shape[1:])
+        segment_key = None if key_rows is None else key_rows[:, start:end].contiguous().reshape(-1, *key.shape[1:])
+        segment_value = (
+            None if value_rows is None else value_rows[:, start:end].contiguous().reshape(-1, *value.shape[1:])
+        )
+        segment_output = segment_runner(segment_query, segment_key, segment_value, metadata, None)
+        expected_tokens = batch_size * (end - start)
+        if segment_output.shape[0] != expected_tokens:
+            raise ValueError(
+                f"Paged piecewise runner returned {segment_output.shape[0]} tokens "
+                f"for a {expected_tokens}-token homogeneous segment"
+            )
+        segment_outputs.append(segment_output.reshape(batch_size, end - start, *segment_output.shape[1:]))
+
+    if not segment_outputs:
+        result = torch.empty_like(query)
+    elif len(segment_outputs) == 1:
+        result = segment_outputs[0].reshape(plan.num_query_tokens, *segment_outputs[0].shape[2:])
+    else:
+        result_rows = torch.cat(segment_outputs, dim=1)
+        result = result_rows.reshape(plan.num_query_tokens, *result_rows.shape[2:])
+    if output_buffer is not None:
+        if output_buffer.shape != result.shape:
+            raise ValueError(
+                f"Paged piecewise output buffer shape {tuple(output_buffer.shape)} does not match "
+                f"homogeneous result shape {tuple(result.shape)}"
+            )
+        output_buffer.copy_(result)
+        return output_buffer
+    return result
 
 
 def run_paged_piecewise_plan(
     query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
+    key: torch.Tensor | None,
+    value: torch.Tensor | None,
     plan: PagedPiecewisePlan,
     segment_metadata: Sequence[object],
     segment_runner: PagedPiecewiseRunner,
+    output_buffer: torch.Tensor | None = None,
+    *,
+    use_homogeneous_batch: bool = False,
 ) -> torch.Tensor:
-    """Run and scatter piecewise attention over packed Q/K/V write tensors."""
+    """Run piecewise attention with view fast paths for contiguous segments.
+
+    ``use_homogeneous_batch`` is an opt-in for native backends that can keep
+    identical rows in the legacy batch layout.  The default remains the
+    indexed packed path so existing GPU/heterogeneous contracts are unchanged.
+    """
 
     if query.shape[0] != plan.num_query_tokens:
         raise ValueError(f"Paged piecewise plan has {plan.num_query_tokens} query tokens, got {query.shape[0]}")
+    if output_buffer is not None and output_buffer.shape[0] != plan.num_query_tokens:
+        raise ValueError(
+            f"Paged piecewise output buffer has {output_buffer.shape[0]} tokens; expected {plan.num_query_tokens}"
+        )
 
-    output = None
+    if use_homogeneous_batch and plan.homogeneous_batch_shape is not None:
+        return _run_homogeneous_paged_piecewise_plan(
+            query,
+            key,
+            value,
+            plan,
+            segment_metadata,
+            segment_runner,
+            output_buffer,
+        )
+
+    output = output_buffer
+    contiguous_outputs = []
     for segment, metadata in zip(plan.segments, segment_metadata, strict=True):
         indices = segment.query_indices
+        segment_output_buffer = None
+        if segment.query_range is None:
+            segment_query = query.index_select(0, indices)
+            segment_key = None if key is None else key.index_select(0, indices)
+            segment_value = None if value is None else value.index_select(0, indices)
+        else:
+            start, end = segment.query_range
+            segment_query = query[start:end]
+            segment_key = None if key is None else key[start:end]
+            segment_value = None if value is None else value[start:end]
+            if output is not None:
+                segment_output_buffer = output[start:end]
         segment_output = segment_runner(
-            query.index_select(0, indices),
-            key.index_select(0, indices),
-            value.index_select(0, indices),
+            segment_query,
+            segment_key,
+            segment_value,
             metadata,
+            segment_output_buffer,
         )
+        if segment_output.shape[0] != indices.shape[0]:
+            raise ValueError(
+                f"Paged piecewise runner returned {segment_output.shape[0]} tokens "
+                f"for a {indices.shape[0]}-token segment"
+            )
+        if segment_output_buffer is not None and segment_output is not segment_output_buffer:
+            if segment_output.shape != segment_output_buffer.shape:
+                raise ValueError(
+                    f"Paged piecewise runner returned shape {tuple(segment_output.shape)} "
+                    f"for output buffer shape {tuple(segment_output_buffer.shape)}"
+                )
+            segment_output_buffer.copy_(segment_output)
+            segment_output = segment_output_buffer
+        if plan.segments_cover_query_contiguously:
+            if output_buffer is None:
+                contiguous_outputs.append(segment_output)
+            continue
         if output is None:
             output = segment_output.new_empty((query.shape[0], *segment_output.shape[1:]))
-        output.index_copy_(0, indices, segment_output)
+        if segment.query_range is None:
+            output.index_copy_(0, indices, segment_output)
+        else:
+            start, end = segment.query_range
+            output[start:end].copy_(segment_output)
 
+    if output_buffer is not None and plan.segments_cover_query_contiguously:
+        return output_buffer
+    if contiguous_outputs:
+        if len(contiguous_outputs) == 1:
+            return contiguous_outputs[0]
+        return torch.cat(contiguous_outputs, dim=0)
     assert output is not None
     return output

vllm_omni/diffusion/attention/layer.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
 # Copyright (c) Microsoft Corporation and Jiarui Fang
 # SPDX-License-Identifier: Apache-2.0
 # DeepSpeed Team & Jiarui Fang
@@ -23,6 +23,7 @@ from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
 from vllm_omni.diffusion.forward_context import (
     get_forward_context,
@@ -44,6 +45,8 @@ def _try_extract_layer_index(prefix: str) -> int | None:
 
 
 class Attention(nn.Module):
+    _scheduler_paged_kv = False
+
     def __init__(
         self,
         num_heads: int,
@@ -111,6 +114,44 @@ class Attention(nn.Module):
             role_category=role_category,
             allow_trtllm_default=allow_trtllm_default,
         )
+        scheduler_paged_kv = (
+            config is not None
+            and getattr(config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            is DiffusionKVCacheMode.PAGED_SCHEDULER
+        )
+        self._scheduler_paged_kv = scheduler_paged_kv
+        if (
+            scheduler_paged_kv
+            and paged_kv_cache_role is not None
+            and spec is None
+            and not attn_backend_cls.supports_paged_kv
+        ):
+            # FLASH_ATTN is an Omni selector, not a device-specific kernel.
+            # vLLM resolves it to CUDA FlashAttention on GPU and Ascend native
+            # attention on NPU. A platform may legitimately default dense
+            # attention to SDPA (for example when an optional diffusion FA
+            # package is absent), but a paged layer must advertise the
+            # capability so the Worker can register its native cache view.
+            # Resolve the selector only for this marked layer; unmarked dense
+            # attention keeps the platform default unchanged.
+            from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+
+            dense_backend_name = attn_backend_cls.get_name()
+            paged_attention_config = AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN"))
+            attn_backend_cls, spec = get_attn_backend_for_role(
+                role=role,
+                head_size=head_size,
+                attention_config=paged_attention_config,
+                role_category=role_category,
+                allow_trtllm_default=allow_trtllm_default,
+            )
+            logger.info(
+                "Resolved marked paged diffusion attention role=%r to %r because the platform default %r "
+                "does not support Scheduler-owned KV",
+                role,
+                attn_backend_cls.get_name(),
+                dense_backend_name,
+            )
         parallel_config = getattr(config, "parallel_config", None)
         allgather_degree = getattr(parallel_config, "allgather_degree", 1)
         # TODO: Move AllGather-KV compatibility into an AttentionBackend capability
@@ -321,6 +362,17 @@ class Attention(nn.Module):
         # Get the appropriate parallel strategy based on SP active state
         strategy = self._get_active_parallel_strategy()
         paged_adapter = self._active_paged_kv_adapter()
+        in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
+        if (
+            self._scheduler_paged_kv
+            and self.paged_kv_cache_role is not None
+            and paged_adapter is None
+            and not in_kv_memory_profile
+        ):
+            raise RuntimeError(
+                "Scheduler-paged diffusion attention reached model forward without an active Worker adapter. "
+                "Only the startup KV memory profile may execute before paged KV initialization."
+            )
         use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None
         if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False):
             raise NotImplementedError(
@@ -343,9 +395,39 @@ class Attention(nn.Module):
         # For Ring: Concat joint_q
         query, key, value, attn_metadata, ctx = strategy.pre_attention(query, key, value, attn_metadata)
 
-        # 2. Kernel execution stays inside the selected Omni backend.  The
-        # Worker adapter only prepares the native page-table context after
-        # SP has produced rank-local Q/K/V tensors.
+        # Scheduler rows describe the logical sequence, while strict Ulysses
+        # may append synthetic tokens solely to make the image shard divisible.
+        # Remove those tokens after the all-to-all and put zero placeholders
+        # back before the reverse all-to-all.
+        paged_sp_padding = 0
+        paged_sp_padding_offset = query.shape[1]
+        if use_paged_attention and strategy is not self._no_parallel_strategy and strategy.name == "ulysses":
+            forward_ctx = get_forward_context() if is_forward_context_available() else None
+            paged_sp_padding = int(getattr(forward_ctx, "sp_padding_size", 0))
+            if paged_sp_padding:
+                joint_len = int(getattr(ctx, "joint_len", 0))
+                joint_strategy = str(getattr(ctx, "joint_strategy", "front"))
+                paged_sp_padding_offset = query.shape[1] - (joint_len if joint_strategy == "rear" else 0)
+                padding_start = paged_sp_padding_offset - paged_sp_padding
+                if padding_start < 0:
+                    raise ValueError(
+                        "Paged Ulysses padding exceeds the post-all-to-all sequence: "
+                        f"padding={paged_sp_padding}, sequence={query.shape[1]}"
+                    )
+
+                def _remove_paged_sp_padding(tensor: torch.Tensor) -> torch.Tensor:
+                    return torch.cat(
+                        (tensor[:, :padding_start], tensor[:, paged_sp_padding_offset:]),
+                        dim=1,
+                    ).contiguous()
+
+                query = _remove_paged_sp_padding(query)
+                key = _remove_paged_sp_padding(key)
+                value = _remove_paged_sp_padding(value)
+
+        # 2. This is the shared GPU/NPU boundary. The Worker adapter prepares
+        # the native page-table context after SP has produced rank-local Q/K/V;
+        # backend resolution below selects CUDA or Ascend execution.
         if use_paged_attention:
             assert paged_adapter is not None
             paged_kv_context = paged_adapter.prepare_layer_context(
@@ -363,6 +445,11 @@ class Attention(nn.Module):
             else:
                 out = self._run_local_attention(query, key, value, attn_metadata)
 
+        if paged_sp_padding:
+            padding_start = paged_sp_padding_offset - paged_sp_padding
+            output_padding = out.new_zeros((out.shape[0], paged_sp_padding, *out.shape[2:]))
+            out = torch.cat((out[:, :padding_start], output_padding, out[:, padding_start:]), dim=1)
+
         # 3. Post-processing (Reverse Communication)
         # For Ulysses: AllToAll Output, and AllGather Joint Output
         out = strategy.post_attention(out, ctx)
@@ -371,7 +458,7 @@ class Attention(nn.Module):
 
     @staticmethod
     def _active_paged_kv_adapter():
-        """Return the opaque Worker adapter installed for this forward."""
+        """Return the Worker adapter selected by Runner-owned metadata."""
 
         if not is_forward_context_available():
             return None
@@ -392,7 +479,25 @@ class Attention(nn.Module):
             )
             return self.sdpa_fallback.forward(query, key, value, attn_metadata)
 
-        # Fallback to standard attention
+        in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
+        # NPU only: the startup KV-capacity profile needs tensor shapes, not a
+        # paged attention result. If dense MindIE-SD is absent, SDPA provides
+        # that profile forward. Formal GPU/NPU paged requests never use this
+        # branch because their Worker adapter is active.
+        if (
+            self._scheduler_paged_kv
+            and self.paged_kv_cache_role is not None
+            and in_kv_memory_profile
+            and current_omni_platform.is_npu()
+            and self.attn_backend.get_name() == "FLASH_ATTN"
+            and not current_omni_platform.supports_diffusion_dense_flash_attention()
+        ):
+            logger.warning_once(
+                "The startup KV memory profile is using SDPA because MindIE-SD is unavailable. "
+                "Formal paged requests still use the platform-native paged attention backend."
+            )
+            return self.sdpa_fallback.forward(query, key, value, attn_metadata)
+
         return self.attention.forward(query, key, value, attn_metadata)
 
     def _assert_piecewise_compatible(self, attn_metadata: AttentionMetadata | None) -> None:

vllm_omni/diffusion/data.py

@@ -1086,10 +1086,9 @@ class OmniDiffusionConfig:
             "need_recv_cache", False
         ):
             raise ValueError(
-                "paged_scheduler Diffusion KV does not support imported AR KV in Phase 1; "
-                "disable need_recv_cache until connector-aware admission is implemented"
+                "paged_scheduler Diffusion KV does not support imported AR KV; "
+                "disable need_recv_cache until connector-aware import is implemented"
             )
-
         self.master_port = self._resolve_master_port()
         self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0)
         if not math.isfinite(self.request_batch_max_wait_ms) or self.request_batch_max_wait_ms < 0:

vllm_omni/diffusion/diffusion_kv/model_runner_backend.py

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
 import copy
-from collections.abc import Mapping, Sequence
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Any
 
@@ -22,9 +22,9 @@ from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionAdapter,
     DiffusionPagedAttentionLayerAdapter,
-    DiffusionPagedAttentionRow,
+    DiffusionPagedAttentionMetadata,
     DiffusionPagedAttentionRowBinding,
-    PreparedDiffusionPagedAttentionBatch,
+    DiffusionPagedAttentionRuntime,
 )
 from vllm_omni.platforms import current_omni_platform
 
@@ -32,13 +32,6 @@ DiffusionKVIdentity = tuple[str, int | None, str | None]
 DiffusionKVSnapshot = tuple[object, ...]
 
 
-@dataclass(frozen=True)
-class _DiffusionKVRequestState:
-    generation: int
-    snapshot: DiffusionKVSnapshot
-    row_token_lens: tuple[tuple[DiffusionKVIdentity, int], ...]
-
-
 @dataclass(frozen=True)
 class _DiffusionKVRowInstall:
     identity: DiffusionKVIdentity
@@ -46,6 +39,13 @@ class _DiffusionKVRowInstall:
     block_ids: tuple[tuple[int, ...], ...]
 
 
+@dataclass(frozen=True)
+class _DiffusionKVRequestState:
+    generation: int
+    snapshot: DiffusionKVSnapshot
+    row_installs: tuple[_DiffusionKVRowInstall, ...]
+
+
 class DiffusionKVModelRunnerBackend:
     """Native paged-KV state and BlockTable operations for a model runner."""
 
@@ -503,7 +503,7 @@ class DiffusionKVModelRunnerBackend:
             self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
                 generation=metadata.allocation_generation,
                 snapshot=snapshot,
-                row_token_lens=current_state.row_token_lens,
+                row_installs=current_state.row_installs,
             )
             return False
 
@@ -530,7 +530,7 @@ class DiffusionKVModelRunnerBackend:
         self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
             generation=metadata.allocation_generation,
             snapshot=snapshot,
-            row_token_lens=tuple((install.identity, install.token_len) for install in installs),
+            row_installs=tuple(installs),
         )
         return True
 
@@ -559,24 +559,11 @@ class DiffusionKVModelRunnerBackend:
             raise RuntimeError("paged_scheduler native attention adapter is not initialized")
         return self.paged_attention_adapter
 
-    def prepare_paged_attention_batch(
-        self,
-        rows: Sequence[DiffusionPagedAttentionRow],
-    ) -> PreparedDiffusionPagedAttentionBatch:
-        """Build native page-table metadata for one Diffusion forward.
+    def activate_paged_attention_metadata(self, metadata: DiffusionPagedAttentionMetadata):
+        """Activate Runner-owned request metadata for an internal denoise loop."""
 
-        The scheduler allocation payload describes capacity, while the model
-        integration supplies the current ``query_len``/``kv_start_pos`` span
-        for each row.  Keeping this boundary explicit prevents the Worker from
-        guessing model-specific text/image layout.
-        """
-
-        return self.get_paged_attention_adapter().prepare_batch(rows)
-
-    def activate_paged_attention(self, batch: PreparedDiffusionPagedAttentionBatch):
-        """Return the context manager that exposes a prepared batch to Omni Attention."""
-
-        return self.get_paged_attention_adapter().activate(batch)
+        runtime = DiffusionPagedAttentionRuntime(self.get_paged_attention_adapter(), metadata)
+        return runtime, runtime.activate()
 
     def _resolve_paged_attention_row(
         self,
@@ -588,11 +575,12 @@ class DiffusionKVModelRunnerBackend:
         identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None)
         request_state = self._diffusion_kv_request_states.get(request_id)
         if request_state is not None:
-            for installed_identity, token_len in request_state.row_token_lens:
-                if installed_identity == identity:
+            for install in request_state.row_installs:
+                if install.identity == identity:
                     return DiffusionPagedAttentionRowBinding(
                         row_index=row_index,
-                        max_seq_len=token_len,
+                        max_seq_len=install.token_len,
+                        block_ids=install.block_ids,
                     )
         raise RuntimeError(f"Diffusion KV request state is missing logical length for {identity!r}")
 

vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.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
 
 from __future__ import annotations
 
@@ -34,6 +34,7 @@ class DiffusionPagedAttentionRowBinding:
 
     row_index: int
     max_seq_len: int
+    block_ids: tuple[tuple[int, ...], ...] = ()
 
 
 DiffusionKVRowResolver = Callable[
@@ -100,6 +101,10 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
         finally:
             attention_config.backend = previous_backend
             attention_config.backend_per_kind = previous_backend_per_kind
+        attn_backend = current_omni_platform.get_diffusion_paged_kv_attn_backend(
+            attn_backend,
+            ulysses_degree=ulysses_degree,
+        )
         canonical_spec = replace(
             spec,
             num_kv_heads=num_kv_heads,
@@ -152,9 +157,11 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
             raise RuntimeError(
                 f"Native paged-attention backend {self.attn_backend.get_name()!r} has no forward implementation"
             )
-        if not self.attn_backend.forward_includes_kv_cache_update and not callable(
-            getattr(impl, "do_kv_cache_update", None)
-        ):
+        needs_explicit_cache_update = (
+            not self.attn_backend.forward_includes_kv_cache_update
+            or current_omni_platform.requires_diffusion_paged_kv_prewrite()
+        )
+        if needs_explicit_cache_update and not callable(getattr(impl, "do_kv_cache_update", None)):
             raise RuntimeError(
                 f"Native paged-attention backend {self.attn_backend.get_name()!r} cannot update the KV cache"
             )
@@ -204,6 +211,30 @@ class DiffusionPagedAttentionRow:
         return (self.request_id, self.sequence_id, self.context_id)
 
 
+@dataclass(frozen=True)
+class DiffusionPagedAttentionMetadata:
+    """Runner-owned row layouts for one request-level denoise loop."""
+
+    prefill_rows: tuple[DiffusionPagedAttentionRow, ...]
+    denoise_rows: tuple[DiffusionPagedAttentionRow, ...]
+
+    def __post_init__(self) -> None:
+        if not self.prefill_rows or not self.denoise_rows:
+            raise ValueError("Paged attention metadata requires both prefill and denoise rows")
+        prefill_identities = tuple(row.identity for row in self.prefill_rows)
+        denoise_identities = tuple(row.identity for row in self.denoise_rows)
+        if prefill_identities != denoise_identities:
+            raise ValueError(
+                "Paged attention prefill and denoise rows must use the same ordered identities: "
+                f"prefill={prefill_identities!r}, denoise={denoise_identities!r}"
+            )
+
+    def rows_for_step(self, step_idx: int | None) -> tuple[DiffusionPagedAttentionRow, ...]:
+        if type(step_idx) is not int or step_idx < 0:
+            raise ValueError(f"Paged attention requires a non-negative denoise step index, got {step_idx!r}")
+        return self.prefill_rows if step_idx == 0 else self.denoise_rows
+
+
 @dataclass(frozen=True)
 class PreparedDiffusionPagedAttentionBatch:
     """Native metadata shared by all paged attention layers in one forward."""
@@ -236,13 +267,81 @@ class DiffusionPagedAttentionContext:
     piecewise_native_metadata: tuple[Any, ...]
     query_token_shape: tuple[int, ...]
     query_has_head_dims: bool
+    output_scatter_indices: torch.Tensor | None = None
 
     def restore_output(self, output: torch.Tensor) -> torch.Tensor:
+        if self.output_scatter_indices is not None:
+            flat_output = output.reshape(-1, self.layer.num_heads, self.layer.head_size_v)
+            restored = output.new_zeros(
+                (*self.query_token_shape, self.layer.num_heads, self.layer.head_size_v)
+            ).reshape(-1, self.layer.num_heads, self.layer.head_size_v)
+            restored.index_copy_(0, self.output_scatter_indices, flat_output)
+            if self.query_has_head_dims:
+                return restored.reshape(
+                    *self.query_token_shape,
+                    self.layer.num_heads,
+                    self.layer.head_size_v,
+                )
+            return restored.reshape(
+                *self.query_token_shape,
+                self.layer.num_heads * self.layer.head_size_v,
+            )
         if self.query_has_head_dims:
             return output.reshape(*self.query_token_shape, self.layer.num_heads, self.layer.head_size_v)
         return output.reshape(*self.query_token_shape, self.layer.num_heads * self.layer.head_size_v)
 
 
+class DiffusionPagedAttentionRuntime:
+    """Activate Runner-prepared rows lazily as the denoise loop advances."""
+
+    def __init__(
+        self,
+        adapter: DiffusionPagedAttentionAdapter,
+        metadata: DiffusionPagedAttentionMetadata,
+    ) -> None:
+        self.adapter = adapter
+        self.metadata = metadata
+        self._active_phase: bool | None = None
+        self._activation: Any | None = None
+        self._closed = False
+
+    def ensure_active(self, step_idx: int | None) -> DiffusionPagedAttentionAdapter:
+        if self._closed:
+            raise RuntimeError("Paged attention runtime is already closed")
+        is_prefill = step_idx == 0
+        rows = self.metadata.rows_for_step(step_idx)
+        if self._active_phase == is_prefill:
+            return self.adapter
+
+        self._deactivate()
+        batch = self.adapter.prepare_batch(rows)
+        activation = self.adapter.activate(batch)
+        activation.__enter__()
+        self._activation = activation
+        self._active_phase = is_prefill
+        return self.adapter
+
+    def _deactivate(self) -> None:
+        activation = self._activation
+        self._activation = None
+        self._active_phase = None
+        if activation is not None:
+            activation.__exit__(None, None, None)
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._deactivate()
+        self._closed = True
+
+    @contextmanager
+    def activate(self) -> Iterator[DiffusionPagedAttentionRuntime]:
+        try:
+            yield self
+        finally:
+            self.close()
+
+
 class DiffusionPagedAttentionAdapter:
     """Translate diffusion rows into vLLM metadata for an Omni paged backend."""
 
@@ -545,15 +644,16 @@ class DiffusionPagedAttentionAdapter:
         )
 
     @staticmethod
-    def _validate_token_layout(
+    def _token_packing_indices(
         token_shape: tuple[int, ...],
         batch: PreparedDiffusionPagedAttentionBatch,
         *,
         name: str,
-    ) -> None:
+        device: torch.device,
+    ) -> torch.Tensor | None:
         if len(token_shape) == 1:
             if token_shape[0] == batch.num_tokens:
-                return
+                return None
             raise ValueError(
                 f"Paged attention {name} token count must match the prepared write batch: "
                 f"tokens={token_shape[0]}, prepared={batch.num_tokens}"
@@ -561,11 +661,23 @@ class DiffusionPagedAttentionAdapter:
         if len(token_shape) == 2:
             batch_size, tokens_per_row = token_shape
             query_lens = tuple(row.query_len for row in batch.rows)
-            if batch_size == len(batch.rows) and all(query_len == tokens_per_row for query_len in query_lens):
-                return
-            raise ValueError(
-                f"Paged attention batched {name} layout must match prepared rows for the current write: "
-                f"shape={token_shape}, row_query_lens={query_lens}"
+            if batch_size != len(batch.rows) or any(query_len > tokens_per_row for query_len in query_lens):
+                raise ValueError(
+                    f"Paged attention batched {name} layout must match prepared rows for the current write: "
+                    f"shape={token_shape}, row_query_lens={query_lens}"
+                )
+            if all(query_len == tokens_per_row for query_len in query_lens):
+                return None
+            return torch.cat(
+                [
+                    torch.arange(
+                        row_index * tokens_per_row,
+                        row_index * tokens_per_row + query_len,
+                        dtype=torch.long,
+                        device=device,
+                    )
+                    for row_index, query_len in enumerate(query_lens)
+                ]
             )
         raise ValueError(
             f"Paged attention {name} supports packed [T, ...] or uniform batched [B, T, ...] token layouts; "
@@ -626,19 +738,39 @@ class DiffusionPagedAttentionAdapter:
                 query_offsets.append(query_offsets[-1] + query_len)
             query_start_loc_cpu = torch.tensor(query_offsets, dtype=torch.int32)
             seq_lens_cpu = torch.tensor(seq_lens_values, dtype=torch.int32)
-            metadata_by_segment.append(
-                self._build_native_metadata(
-                    query_lens=query_lens,
-                    query_start_loc=query_start_loc_cpu.to(self.device),
-                    query_start_loc_cpu=query_start_loc_cpu,
-                    seq_lens=seq_lens_cpu.to(self.device),
-                    seq_lens_cpu=seq_lens_cpu,
-                    positions=batch.positions.index_select(0, packed_segment.query_indices),
-                    block_tables=batch.block_tables,
-                    slot_mappings=batch.slot_mappings.index_select(-1, packed_segment.query_indices),
-                    causal=(row_segments[0].mode == "causal"),
-                )
+            if packed_segment.query_range is None:
+                positions = batch.positions.index_select(0, packed_segment.query_indices)
+                slot_mappings = batch.slot_mappings.index_select(-1, packed_segment.query_indices)
+            else:
+                start, end = packed_segment.query_range
+                positions = batch.positions[start:end]
+                slot_mappings = batch.slot_mappings[..., start:end]
+            segment_metadata = self._build_native_metadata(
+                query_lens=query_lens,
+                query_start_loc=query_start_loc_cpu.to(self.device),
+                query_start_loc_cpu=query_start_loc_cpu,
+                seq_lens=seq_lens_cpu.to(self.device),
+                seq_lens_cpu=seq_lens_cpu,
+                positions=positions,
+                block_tables=batch.block_tables,
+                slot_mappings=slot_mappings,
+                causal=(row_segments[0].mode == "causal"),
             )
+            # FA3's full-CUDA-graph metadata builder reuses one persistent
+            # scheduler buffer across builds. Piecewise attention prepares all
+            # segments before executing any of them, so each segment needs its
+            # own snapshot rather than a view overwritten by the next build.
+            snapshotted = set()
+            for native_metadata in segment_metadata.values():
+                metadata_id = id(native_metadata)
+                if metadata_id in snapshotted:
+                    continue
+                snapshotted.add(metadata_id)
+                for field_name in ("scheduler_metadata", "prefix_scheduler_metadata"):
+                    scheduler_metadata = getattr(native_metadata, field_name, None)
+                    if isinstance(scheduler_metadata, torch.Tensor):
+                        setattr(native_metadata, field_name, scheduler_metadata.clone())
+            metadata_by_segment.append(segment_metadata)
         self._active_piecewise_native_metadata = tuple(metadata_by_segment)
         return self._active_piecewise_native_metadata
 
@@ -678,9 +810,36 @@ class DiffusionPagedAttentionAdapter:
             head_size=layer.head_size_v,
             name="value",
         )
-        self._validate_token_layout(query_token_shape, batch, name="query")
-        self._validate_token_layout(key_token_shape, batch, name="key")
-        self._validate_token_layout(value_token_shape, batch, name="value")
+        if key_token_shape != query_token_shape:
+            if len(key_token_shape) == 1:
+                raise ValueError(
+                    "Paged attention key token count must match the query layout: "
+                    f"query={query_token_shape}, key={key_token_shape}"
+                )
+            raise ValueError(
+                "Paged attention batched key layout must match prepared rows and the query layout: "
+                f"query={query_token_shape}, key={key_token_shape}"
+            )
+        if value_token_shape != query_token_shape:
+            if len(value_token_shape) == 1:
+                raise ValueError(
+                    "Paged attention value token count must match the query layout: "
+                    f"query={query_token_shape}, value={value_token_shape}"
+                )
+            raise ValueError(
+                "Paged attention batched value layout must match prepared rows and the query layout: "
+                f"query={query_token_shape}, value={value_token_shape}"
+            )
+        output_scatter_indices = self._token_packing_indices(
+            query_token_shape,
+            batch,
+            name="Q/K/V",
+            device=query.device,
+        )
+        if output_scatter_indices is not None:
+            query_flat = query_flat.index_select(0, output_scatter_indices)
+            key_flat = key_flat.index_select(0, output_scatter_indices)
+            value_flat = value_flat.index_select(0, output_scatter_indices)
         if query.device != self.device or key.device != self.device or value.device != self.device:
             raise ValueError(
                 f"Paged attention Q/K/V must be on {self.device}; "
@@ -711,10 +870,18 @@ class DiffusionPagedAttentionAdapter:
                 f"expected {batch.num_tokens}"
             )
 
-        # vLLM's writer consumes slot_mapping directly, so a span that starts
+        # vLLM's CUDA writer consumes the source token strides directly, so
+        # preserve strided QKV-projection views instead of materializing an
+        # extra K/V copy here. Platform writers with stricter contracts (for
+        # example, Ascend PA scatter) make their inputs contiguous locally.
+        # The native CUDA writer still assumes heads and head dimensions are
+        # packed within each token; materialize unusual inner layouts only.
+        if key_flat.stride(-1) != 1 or (key_flat.shape[-2] > 1 and key_flat.stride(-2) != key_flat.shape[-1]):
+            key_flat = key_flat.contiguous()
+        if value_flat.stride(-1) != 1 or (value_flat.shape[-2] > 1 and value_flat.stride(-2) != value_flat.shape[-1]):
+            value_flat = value_flat.contiguous()
+        # The writer also consumes slot_mapping directly, so a span that starts
         # inside a physical block naturally writes only the requested suffix.
-        key_flat = key_flat.contiguous()
-        value_flat = value_flat.contiguous()
         try:
             native_metadata = batch.attn_metadata[layer_name]
         except KeyError as exc:
@@ -743,6 +910,7 @@ class DiffusionPagedAttentionAdapter:
             piecewise_native_metadata=piecewise_native_metadata,
             query_token_shape=query_token_shape,
             query_has_head_dims=query_has_head_dims,
+            output_scatter_indices=output_scatter_indices,
         )
 
     @staticmethod

vllm_omni/diffusion/forward_context.py

@@ -29,11 +29,19 @@ class ForwardContext:
     vllm_config: VllmConfig | None = None
     omni_diffusion_config: OmniDiffusionConfig | None = None
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None
+    # Runner-owned paged execution metadata/runtime. Attention resolves the
+    # active Worker adapter from it; model code must not construct BlockTable
+    # rows or activate the runtime directly.
+    paged_kv_runtime: object | None = None
     # Active Worker-side paged KV adapter.  The adapter is installed only for
     # the duration of a paged forward; dense forwards leave this as ``None``.
     # Keep the field opaque here to avoid coupling the common context module to
     # the diffusion_kv implementation.
     paged_kv_adapter: Any | None = None
+    # Startup-only memory profiling runs before Scheduler-owned pages exist.
+    # Attention layers use this explicit scope to distinguish that probe from
+    # a malformed paged request whose Worker adapter was not activated.
+    in_diffusion_kv_memory_profile: bool = False
     split_text_embed_in_sp: bool = False
     denoise_step_idx: int | None = None
     denoise_timestep: float | None = None
@@ -154,6 +162,8 @@ def create_forward_context(
     vllm_config: VllmConfig | None = None,
     omni_diffusion_config: OmniDiffusionConfig | None = None,
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
+    paged_kv_runtime: object | None = None,
+    in_diffusion_kv_memory_profile: bool = False,
     split_text_embed_in_sp: bool = False,
     denoise_step_idx: int | None = None,
 ):
@@ -161,6 +171,8 @@ def create_forward_context(
         vllm_config=vllm_config,
         omni_diffusion_config=omni_diffusion_config,
         attn_metadata=attn_metadata,
+        paged_kv_runtime=paged_kv_runtime,
+        in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
         split_text_embed_in_sp=split_text_embed_in_sp,
         denoise_step_idx=denoise_step_idx,
     )
@@ -186,6 +198,8 @@ def set_forward_context(
     vllm_config: VllmConfig | None = None,
     omni_diffusion_config: OmniDiffusionConfig | None = None,
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
+    paged_kv_runtime: object | None = None,
+    in_diffusion_kv_memory_profile: bool = False,
     split_text_embed_in_sp: bool = False,
     denoise_step_idx: int | None = None,
 ):
@@ -197,6 +211,8 @@ def set_forward_context(
         vllm_config=vllm_config,
         omni_diffusion_config=omni_diffusion_config,
         attn_metadata=attn_metadata,
+        paged_kv_runtime=paged_kv_runtime,
+        in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
         split_text_embed_in_sp=split_text_embed_in_sp,
         denoise_step_idx=denoise_step_idx,
     )
@@ -246,6 +262,11 @@ def set_forward_context_denoise_step_idx(step_idx: int | None) -> None:
     """Set the current diffusion denoise step on the active ForwardContext."""
     if _forward_context is not None:
         _forward_context.denoise_step_idx = step_idx
+        if step_idx is not None:
+            paged_kv_runtime = getattr(_forward_context, "paged_kv_runtime", None)
+            ensure_active = getattr(paged_kv_runtime, "ensure_active", None)
+            if callable(ensure_active):
+                ensure_active(step_idx)
 
 
 def set_forward_context_denoise_timestep(timestep: float | None) -> None:

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.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
 
 import inspect
 import logging
@@ -1069,18 +1069,98 @@ class ImageKVCacheManager(nn.Module):
         )
         return key, value
 
-    def forward(
+    def clear_legacy_prompt_kv_cache(self) -> None:
+        """Drop model-owned dense KV state before Scheduler-paged execution."""
+
+        self.image_kv_cache_map = None
+        self.image_kv_cache_lens = None
+
+    def _forward_paged(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
+        *,
+        first_step: bool,
+        query_lens: list[int],
+        seq_lens: list[int],
+        shard_image_size: int | None,
+        full_attn_spans: list[list[tuple[int, int]]],
+        uncond_cfg_prefill: bool,
+    ) -> torch.Tensor:
+        """Run Hunyuan attention against Worker-managed Scheduler pages.
+
+        This model boundary is shared by GPU and NPU: it describes Hunyuan's
+        Q/K/V layout and mixed-attention spans, while the common Attention
+        layer resolves the actual CUDA or Ascend paged kernel.
+        """
+
+        if uncond_cfg_prefill:
+            raise RuntimeError("Hunyuan negative-CFG prefill must run before paged row activation")
+        if self._injected_ar_kv is not None:
+            raise NotImplementedError("Hunyuan Scheduler-paged KV does not support imported AR KV")
+        if not query_lens or len(seq_lens) != len(query_lens):
+            raise ValueError("Hunyuan Scheduler-paged KV requires matching query and sequence lengths")
+        if len(full_attn_spans) != len(query_lens):
+            raise ValueError("Hunyuan Scheduler-paged KV requires one full-attention span row per sequence")
+
+        self.clear_legacy_prompt_kv_cache()
+        bs = len(query_lens)
+        q_len = query_lens[0]
+        seq_len = seq_lens[0]
+        assert query.shape[0] == bs * q_len, f"{query.shape[0]} != {bs * q_len}"
+
+        head_num_per_rank = query.shape[1]
+        kv_head_num_per_rank = key.shape[1]
+        head_dim = query.shape[2]
+
+        query = query.reshape(bs, q_len, head_num_per_rank, head_dim)
+        key = key.reshape(bs, q_len, kv_head_num_per_rank, head_dim)
+        value = value.reshape(bs, q_len, kv_head_num_per_rank, head_dim)
+
+        joint_text_query = joint_text_key = joint_text_value = None
+        if self.sp_size > 1 and first_step:
+            if shard_image_size is None or shard_image_size <= 0:
+                raise ValueError("Hunyuan paged Ulysses requires a positive local image shard size")
+            local_prompt_len = seq_len - shard_image_size
+            joint_query_len = query.shape[1] - shard_image_size
+            if local_prompt_len != joint_query_len:
+                raise ValueError(
+                    "Hunyuan paged Ulysses prompt layout mismatch: "
+                    f"key_prompt={local_prompt_len}, query_prompt={joint_query_len}"
+                )
+            joint_text_query = query[:, :joint_query_len]
+            joint_text_key = key[:, :local_prompt_len]
+            joint_text_value = value[:, :local_prompt_len]
+            query = query[:, joint_query_len:]
+            key = key[:, local_prompt_len:]
+            value = value[:, local_prompt_len:]
+
+        if joint_text_query is None:
+            attn_metadata = AttentionMetadata(full_attn_spans=full_attn_spans)
+        else:
+            attn_metadata = AttentionMetadata(
+                joint_query=joint_text_query,
+                joint_key=joint_text_key,
+                joint_value=joint_text_value,
+                joint_strategy="front",
+                full_attn_spans=full_attn_spans,
+            )
+        attn_output = self.attn(query, key, value, attn_metadata)
+        return attn_output.reshape(bs * q_len, head_num_per_rank, head_dim)
+
+    def _forward_dense_legacy(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: torch.Tensor,
         **kwargs,
     ) -> torch.Tensor:
-        self.image_token_len = kwargs.get("num_image_tokens")
+        """Run the device-neutral, model-owned dense prompt-KV compatibility path."""
+
         first_step = kwargs.get("first_step")
         uncond_cfg_prefill = kwargs.get("uncond_cfg_prefill", False)
-
         query_lens = kwargs.get("query_lens")
         seq_lens = kwargs.get("seq_lens")
         shard_image_size = kwargs.get("shard_image_size") if self.sp_size > 1 else None
@@ -1109,8 +1189,7 @@ class ImageKVCacheManager(nn.Module):
                 key = key[:, :0, :, :]
                 value = value[:, :0, :, :]
         elif first_step:
-            self.image_kv_cache_map = None  # reset first
-            self.image_kv_cache_lens = None
+            self.clear_legacy_prompt_kv_cache()
             key, value = self._cache_prompt_kv(
                 key,
                 value,
@@ -1134,7 +1213,7 @@ class ImageKVCacheManager(nn.Module):
                 joint_text_query = query[:, :0, :, :]
                 joint_text_key, joint_text_value = self._reuse_prompt_kv(key, value, seq_len, bs, shard_image_size)
 
-        if not keep_kv_compressed and not self.attn.is_paged_kv_active():
+        if not keep_kv_compressed:
             key = repeat_kv(key, repeat_num)
             value = repeat_kv(value, repeat_num)
             if self.sp_size > 1:
@@ -1163,6 +1242,34 @@ class ImageKVCacheManager(nn.Module):
         attn_output = attn_output.reshape(bs * q_len, head_num_per_rank, head_dim)
         return attn_output
 
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        self.image_token_len = kwargs.get("num_image_tokens")
+        if self.attn.is_paged_kv_active():
+            full_attn_spans = kwargs.get("full_attn_spans")
+            if full_attn_spans is None:
+                raise ValueError("Hunyuan Scheduler-paged KV requires full_attn_spans metadata")
+            return self._forward_paged(
+                query,
+                key,
+                value,
+                first_step=bool(kwargs.get("first_step")),
+                query_lens=kwargs.get("query_lens"),
+                seq_lens=kwargs.get("seq_lens"),
+                shard_image_size=kwargs.get("shard_image_size") if self.sp_size > 1 else None,
+                full_attn_spans=full_attn_spans,
+                uncond_cfg_prefill=kwargs.get("uncond_cfg_prefill", False),
+            )
+        if attention_mask is None:
+            raise ValueError("Hunyuan dense attention requires an attention mask")
+        return self._forward_dense_legacy(query, key, value, attention_mask, **kwargs)
+
 
 @dataclass
 class CausalMMOutputWithPast(CausalLMOutputWithPast):
@@ -2706,7 +2813,6 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         # List[List[...]] per-sample metadata indexed along the CFG batch dim
         if isinstance(model_kwargs.get("full_attn_spans"), list):
             model_kwargs["full_attn_spans"] = model_kwargs["full_attn_spans"][s.start : s.stop]
-
         # custom_pos_emb: tuple of (cos, sin)
         if "custom_pos_emb" in model_kwargs and model_kwargs["custom_pos_emb"] is not None:
             cos, sin = model_kwargs["custom_pos_emb"]

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.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
 
 import copy
 import logging
@@ -72,7 +72,6 @@ _STEP_OUTPUT_SIZE = "hunyuan_output_size"
 _STEP_COT_TEXT_LIST = "hunyuan_cot_text_list"
 _STEP_AR_KV = "hunyuan_ar_kv"
 _STEP_PROMPT_KV = "hunyuan_prompt_kv"
-
 _HUNYUAN_DEFAULT_OUTPUT_TYPE = "pil"
 
 
@@ -466,6 +465,11 @@ class HunyuanImage3Pipeline(
         return self._pipeline
 
     def _validate_step_request(self, state: "StepRequestState") -> None:
+        if self._uses_scheduler_paged_kv():
+            raise ValueError(
+                "HunyuanImage3 paged_scheduler currently supports request-level execution only; "
+                "disable step execution or use dense_legacy."
+            )
         prompt = state.prompt
         sampling = state.sampling
         if prompt is None:
@@ -663,6 +667,12 @@ class HunyuanImage3Pipeline(
             prefix_lens.append(int(prompt_kv[0]["lens"][branch].item()))
         return prefix_lens
 
+    def _uses_scheduler_paged_kv(self) -> bool:
+        return (
+            getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            is DiffusionKVCacheMode.PAGED_SCHEDULER
+        )
+
     def _merge_step_model_inputs(
         self,
         states: list["StepRequestState"],
@@ -693,6 +703,7 @@ class HunyuanImage3Pipeline(
                         values,
                         row_branches,
                         prefix_lens or [None] * len(row_branches),
+                        strict=True,
                     )
                 ]
                 continue
@@ -1871,8 +1882,7 @@ class HunyuanImage3Pipeline(
         row_branches = [branch for branch in range(cfg_factor) for _ in states]
         return row_state_indexes, row_branches
 
-    @staticmethod
-    def _validate_step_group_states(states: list["StepRequestState"]) -> tuple[bool, int]:
+    def _validate_step_group_states(self, states: list["StepRequestState"]) -> tuple[bool, int]:
         if not states:
             raise ValueError("HunyuanImage3 denoise_step received an empty group.")
 
@@ -2035,6 +2045,11 @@ class HunyuanImage3Pipeline(
         states = list(input_batch.states)
         if not states:
             raise ValueError("HunyuanImage3 denoise_step received an empty batch.")
+        if self._uses_scheduler_paged_kv():
+            raise ValueError(
+                "HunyuanImage3 paged_scheduler currently supports request-level execution only; "
+                "disable step execution or use dense_legacy."
+            )
         self._ensure_grouped_attention_backend_supported(len(states))
         outputs: dict[str, torch.Tensor] = {}
         for group in self._split_step_groups(states):

vllm_omni/diffusion/vllm_config.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
 
 """Build native vLLM configuration views for the Omni diffusion runtime."""
 
@@ -14,6 +14,7 @@ from vllm.transformers_utils.config import get_hf_text_config
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.platforms import current_omni_platform
 
 
 def resolve_diffusion_max_model_len(od_config: OmniDiffusionConfig) -> int:
@@ -227,6 +228,7 @@ def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiff
         max_num_batched_tokens = getattr(od_config, "max_num_batched_tokens", None)
         if max_num_batched_tokens is not None:
             vllm_config.scheduler_config.max_num_batched_tokens = int(max_num_batched_tokens)
+        current_omni_platform.configure_diffusion_vllm_config(vllm_config, od_config)
     return vllm_config
 
 

vllm_omni/diffusion/worker/diffusion_model_runner.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
 """
 Diffusion Model Runner for vLLM-Omni.
 
@@ -37,9 +37,10 @@ from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
+    DiffusionPagedAttentionMetadata,
     DiffusionPagedAttentionRow,
-    PreparedDiffusionPagedAttentionBatch,
 )
+from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_rank
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import (
@@ -420,11 +421,19 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             raise RuntimeError("Model must be loaded before collecting Diffusion KV cache specs")
 
         cache_layers: dict[str, tuple[Attention, KVCacheSpec]] = {}
-        for layer_name, module in self.pipeline.named_modules():
+        for module_path, module in self.pipeline.named_modules():
             if not isinstance(module, Attention):
                 continue
             spec = module.get_kv_cache_spec(self.vllm_config)
             if spec is not None:
+                layer_name = module.prefix
+                if not isinstance(layer_name, str) or not layer_name:
+                    raise RuntimeError(
+                        "Paged Diffusion Attention must expose a non-empty canonical prefix; "
+                        f"module_path={module_path!r}"
+                    )
+                if layer_name in cache_layers:
+                    raise RuntimeError(f"Duplicate canonical paged Diffusion Attention prefix {layer_name!r}")
                 cache_layers[layer_name] = (module, spec)
         if not cache_layers:
             raise RuntimeError(
@@ -456,18 +465,51 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
     def refresh_diffusion_kv_block_table_layout(self) -> None:
         self.diffusion_kv_backend.refresh_block_table_layout()
 
-    def prepare_paged_attention_batch(
+    def _build_paged_attention_metadata(
         self,
-        rows: list[DiffusionPagedAttentionRow] | tuple[DiffusionPagedAttentionRow, ...],
-    ) -> PreparedDiffusionPagedAttentionBatch:
-        """Prepare model-specific query/write spans against installed Worker rows."""
+        metadata: list[DiffusionKVMetadata],
+    ) -> DiffusionPagedAttentionMetadata:
+        """Translate Scheduler sequence state into ordered request-level rows."""
 
-        return self.diffusion_kv_backend.prepare_paged_attention_batch(rows)
-
-    def activate_paged_attention(self, batch: PreparedDiffusionPagedAttentionBatch):
-        """Expose a prepared paged batch to Omni Attention for one forward."""
-
-        return self.diffusion_kv_backend.activate_paged_attention(batch)
+        cfg_size = int(getattr(self.od_config.parallel_config, "cfg_parallel_size", 1) or 1)
+        cfg_rank = get_classifier_free_guidance_rank() if cfg_size > 1 else None
+        prefill_rows: list[DiffusionPagedAttentionRow] = []
+        denoise_rows: list[DiffusionPagedAttentionRow] = []
+        for request_metadata in metadata:
+            sequences = request_metadata.sequences
+            if cfg_rank is not None and len(sequences) > 1:
+                if len(sequences) != cfg_size:
+                    raise ValueError(
+                        "Paged CFG parallel execution requires one Scheduler sequence per CFG rank: "
+                        f"request={request_metadata.request_id!r}, cfg_size={cfg_size}, rows={len(sequences)}"
+                    )
+                sequences = (sequences[cfg_rank],)
+            for sequence in sequences:
+                active_seq_len = sequence.prefix_len + sequence.target_len
+                if active_seq_len > sequence.seq_len:
+                    raise ValueError(
+                        "Paged denoise span exceeds its Scheduler allocation: "
+                        f"request={request_metadata.request_id!r}, sequence={sequence.sequence_id}, "
+                        f"active={active_seq_len}, allocated={sequence.seq_len}"
+                    )
+                prefill_rows.append(
+                    DiffusionPagedAttentionRow(
+                        request_id=request_metadata.request_id,
+                        sequence_id=sequence.sequence_id,
+                        query_len=sequence.seq_len,
+                        seq_len=sequence.seq_len,
+                    )
+                )
+                denoise_rows.append(
+                    DiffusionPagedAttentionRow(
+                        request_id=request_metadata.request_id,
+                        sequence_id=sequence.sequence_id,
+                        query_len=sequence.target_len,
+                        seq_len=active_seq_len,
+                        kv_start_pos=sequence.prefix_len,
+                    )
+                )
+        return DiffusionPagedAttentionMetadata(tuple(prefill_rows), tuple(denoise_rows))
 
     def clear_prompt_embed_cache(self) -> None:
         """Evict all cached text-encoder outputs (e.g. between training epochs).
@@ -623,6 +665,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         kv_prefetch_job: KVPrefetchJob | None = None,
         record_name: str,
         record_output_peak_memory: bool = True,
+        in_diffusion_kv_memory_profile: bool = False,
+        diffusion_kv_metadata: list[DiffusionKVMetadata] | None = None,
     ) -> BatchRunnerOutput:
         assert self.pipeline is not None, "Model not loaded. Call load_model() first."
         if not reqs:
@@ -655,7 +699,27 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             if is_primary and record_output_peak_memory:
                 current_omni_platform.reset_peak_memory_stats()
 
-            with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=od_config):
+            paged_kv_runtime = None
+            paged_kv_context: AbstractContextManager[Any] = nullcontext()
+            if diffusion_kv_metadata is not None:
+                if len(diffusion_kv_metadata) != len(reqs):
+                    raise ValueError(
+                        "Diffusion KV metadata count must match the request batch: "
+                        f"metadata={len(diffusion_kv_metadata)}, requests={len(reqs)}"
+                    )
+                paged_metadata = self._build_paged_attention_metadata(diffusion_kv_metadata)
+                paged_kv_runtime, paged_kv_context = self.diffusion_kv_backend.activate_paged_attention_metadata(
+                    paged_metadata
+                )
+            with (
+                set_forward_context(
+                    vllm_config=self.vllm_config,
+                    omni_diffusion_config=od_config,
+                    paged_kv_runtime=paged_kv_runtime,
+                    in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
+                ),
+                paged_kv_context,
+            ):
                 with record_function(record_name):
                     raw_outputs = self.pipeline.forward(batch)
                     outputs = _normalize_pipeline_outputs(
@@ -733,6 +797,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 require_request_batch_support=False,
                 kv_prefetch_job=kv_prefetch_job,
                 record_name="pipeline_forward",
+                diffusion_kv_metadata=[diffusion_kv_metadata] if diffusion_kv_metadata is not None else None,
             )
             output = runner_output.runner_outputs[0].result
             assert output is not None
@@ -773,6 +838,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     scheduler_output,
                     validate_kv_metadata=False,
                     record_output_peak_memory=False,
+                    in_diffusion_kv_memory_profile=True,
                 )
             else:
                 runner_output = self._execute_request_list(
@@ -785,6 +851,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     # peak counters. Resetting them here would discard request
                     # preparation allocations and understate the budget.
                     record_output_peak_memory=False,
+                    in_diffusion_kv_memory_profile=True,
                 )
             current_omni_platform.synchronize()
         finally:
@@ -824,6 +891,12 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 allow_single_output=False,
                 require_request_batch_support=True,
                 record_name="pipeline_forward_batch",
+                diffusion_kv_metadata=[
+                    new_req.diffusion_kv_metadata
+                    for new_req in scheduler_output.scheduled_new_reqs
+                    if new_req.diffusion_kv_metadata is not None
+                ]
+                or None,
             )
         finally:
             if installed_request_ids:
@@ -984,6 +1057,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         *,
         validate_kv_metadata: bool,
         record_output_peak_memory: bool,
+        in_diffusion_kv_memory_profile: bool = False,
     ) -> BatchRunnerOutput:
         """Execute one step with explicit validation and profiling policy."""
 
@@ -1023,6 +1097,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             return self._execute_stepwise_core(
                 scheduler_output,
                 record_output_peak_memory=record_output_peak_memory,
+                in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
             )
         except Exception:
             if installed_request_ids:
@@ -1034,6 +1109,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         scheduler_output: DiffusionSchedulerOutput,
         *,
         record_output_peak_memory: bool,
+        in_diffusion_kv_memory_profile: bool = False,
     ) -> BatchRunnerOutput:
         """Run the denoise step after metadata admission and row installation."""
 
@@ -1056,10 +1132,14 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 return BatchRunnerOutput.from_list(runner_output_list)
             attn_metadata = {}
 
+            kv_backend = getattr(self, "diffusion_kv_backend", None)
+            paged_kv_runtime = kv_backend if getattr(kv_backend, "paged_attention_adapter", None) is not None else None
             with set_forward_context(
                 vllm_config=self.vllm_config,
                 omni_diffusion_config=self.od_config,
                 attn_metadata=attn_metadata,
+                paged_kv_runtime=paged_kv_runtime,
+                in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
             ):
                 clear_pipeline_stage_durations(self.pipeline)
                 noise_pred = self.pipeline.denoise_step(input_batch, states=states)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -249,11 +249,6 @@ class DiffusionWorker:
         # requests, which only carry their request_id in subsequent ticks.
         self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
         self.stage_id = getattr(od_config, "stage_id", 0)
-        if self.od_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER:
-            logger.warning_once(
-                "paged_scheduler initializes native paged KV storage, but no production diffusion model uses the "
-                "paged-attention adapter yet; model attention remains on the dense path."
-            )
         self.init_device()
         # Create model runner — one decision chain, in precedence order:
         #   1. explicit od_config.diffusion_model_runner_cls (user override),

vllm_omni/platforms/interface.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
 
 from contextlib import nullcontext
 from enum import Enum
@@ -150,6 +150,12 @@ class OmniPlatform(Platform):
             )
         backend_cls.validate_available()
 
+    @classmethod
+    def supports_diffusion_dense_flash_attention(cls) -> bool:
+        """Whether the platform's dense ``FLASH_ATTN`` dependencies exist."""
+
+        return True
+
     @classmethod
     def supports_torch_inductor(cls) -> bool:
         """Check if the platform supports torch.compile with inductor backend."""
@@ -195,17 +201,36 @@ class OmniPlatform(Platform):
     def build_diffusion_kv_attn_metadata(cls, **kwargs: Any) -> dict[str, Any]:
         """Build native attention metadata for the diffusion paged path.
 
-        The common CUDA/ROCm/XPU path uses vLLM's builder. Platforms with a
-        backend-specific metadata contract can override this hook without
-        making the diffusion adapter import their optional runtime package.
-        ``seq_lens_cpu`` is an adapter-only convenience value and is removed
-        before calling the upstream builder.
+        This default is the GPU/common path and uses vLLM's builder. NPU
+        overrides it to build Ascend attention metadata without making the
+        shared diffusion adapter import ``vllm_ascend``. ``seq_lens_cpu`` is an
+        adapter-only convenience value and is removed before calling the
+        upstream builder.
         """
         from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
 
         kwargs.pop("seq_lens_cpu", None)
         return build_attn_metadata(**kwargs)
 
+    @classmethod
+    def get_diffusion_paged_kv_attn_backend(cls, attn_backend: type, *, ulysses_degree: int) -> type:
+        """Specialize a native paged backend for diffusion execution."""
+
+        del ulysses_degree
+        return attn_backend
+
+    @classmethod
+    def requires_diffusion_paged_kv_prewrite(cls) -> bool:
+        """Whether paged attention must write K/V before native execution.
+
+        The default GPU path keeps cache-update ownership in its native
+        attention call. Ascend overrides this because piecewise FIA should
+        write the complete K/V span once, then read it from cache for every
+        segment.
+        """
+
+        return False
+
     @classmethod
     def init_diffusion_worker_vllm_config(
         cls,
@@ -214,6 +239,15 @@ class OmniPlatform(Platform):
         """Initialize platform-specific state for diffusion worker VllmConfig."""
         return None
 
+    @classmethod
+    def configure_diffusion_vllm_config(
+        cls,
+        vllm_config: Any,
+        od_config: Any,
+    ) -> None:
+        """Apply platform-specific native cache geometry for diffusion."""
+        return None
+
     @classmethod
     def init_diffusion_model_runner_runtime(
         cls,

vllm_omni/platforms/npu/platform.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from contextlib import nullcontext
+from functools import cache
 from typing import Any
 
 import torch
@@ -21,6 +22,28 @@ _DIFFUSION_PACKED_MODULES_MAPPING = {
 }
 
 
+@cache
+def _get_strict_ulysses_paged_backend() -> type:
+    """Return an Ascend backend that bypasses vLLM PCP dispatch."""
+
+    from vllm_ascend.attention.attention_v1 import (
+        AscendAttentionBackend,
+        AscendAttentionBackendImpl,
+        AscendAttentionMetadataBuilder,
+    )
+
+    class AscendStrictUlyssesPagedBackend(AscendAttentionBackend):
+        @staticmethod
+        def get_impl_cls() -> type:
+            return AscendAttentionBackendImpl
+
+        @staticmethod
+        def get_builder_cls() -> type:
+            return AscendAttentionMetadataBuilder
+
+    return AscendStrictUlyssesPagedBackend
+
+
 class NPUOmniPlatform(OmniPlatform, NPUPlatform):
     """NPU/Ascend implementation of OmniPlatform.
 
@@ -77,9 +100,55 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
     @classmethod
     def init_diffusion_worker_vllm_config(cls, vllm_config: Any) -> None:
         from vllm_ascend.ascend_config import init_ascend_config
+        from vllm_ascend.utils import adapt_patch
 
+        # Omni's custom DiffusionWorker does not pass through vLLM-Ascend's
+        # NPUWorker constructor, where worker-local patches are normally
+        # installed.  In particular, AscendBlockTables needs the patched
+        # non-UVA buffer implementation on NPU.
+        adapt_patch()
         init_ascend_config(vllm_config)
 
+    @classmethod
+    def configure_diffusion_vllm_config(cls, vllm_config: Any, od_config: Any) -> None:
+        """Use the block geometry required by Ascend's native paged kernel."""
+        if getattr(od_config, "diffusion_kv_mode", None) is None:
+            return
+        from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+
+        if od_config.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
+            return
+        from vllm_ascend.attention.attention_v1 import AscendAttentionBackend
+
+        supported_sizes = [
+            size for size in AscendAttentionBackend.get_supported_kernel_block_sizes() if type(size) is int and size > 0
+        ]
+        if not supported_sizes:
+            raise RuntimeError("Ascend paged attention did not expose an integer kernel block size")
+        # vLLM's generic default is 16, while the Ascend FIA backend stores
+        # cache pages as 128-token blocks. Set the Manager geometry
+        # before KV specs are collected so Scheduler and Worker agree.
+        vllm_config.cache_config.block_size = supported_sizes[0]
+
+    @classmethod
+    def requires_diffusion_paged_kv_prewrite(cls) -> bool:
+        """Write the full K/V span once before piecewise FIA segments."""
+
+        return True
+
+    @classmethod
+    def get_diffusion_paged_kv_attn_backend(cls, attn_backend: type, *, ulysses_degree: int) -> type:
+        """Keep strict Ulysses paged FIA out of vLLM's PCP implementation."""
+
+        del cls
+        if ulysses_degree <= 1:
+            return attn_backend
+        from vllm_ascend.attention.attention_v1 import AscendAttentionBackend
+
+        if not isinstance(attn_backend, type) or not issubclass(attn_backend, AscendAttentionBackend):
+            return attn_backend
+        return _get_strict_ulysses_paged_backend()
+
     @classmethod
     def get_diffusion_kv_block_tables_cls(cls) -> type:
         from vllm_ascend.worker.v2.block_table import AscendBlockTables
@@ -211,6 +280,14 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
         logger.debug("Falling back to diffusion attention backend SDPA")
         return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
 
+    @classmethod
+    def supports_diffusion_dense_flash_attention(cls) -> bool:
+        """Return whether MindIE-SD is installed for dense NPU FlashAttention."""
+
+        from importlib.util import find_spec
+
+        return find_spec("mindiesd") is not None
+
     @classmethod
     def supports_torch_inductor(cls) -> bool:
         return False
@@ -245,7 +322,12 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
         """
         try:
             torch.npu.current_stream().synchronize()
-            event = torch.npu.Event()
+            # The async output worker uses the public ``torch.Stream`` API.
+            # With torch_npu 2.10, a native ``torch.npu.Event`` cannot be
+            # consumed by that wrapper stream (the reverse direction is
+            # supported), while ``torch.Event`` is dispatched to the active
+            # NPU backend and remains cross-stream compatible.
+            event = torch.Event()
             event.record()
             return event
         except Exception: