diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 247f80677..54ad2a32b 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/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): diff --git a/tests/diffusion/attention/test_piecewise_attn.py b/tests/diffusion/attention/test_piecewise_attn.py index d1c414390..fc4746773 100644 --- a/tests/diffusion/attention/test_piecewise_attn.py +++ b/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)]], diff --git a/tests/diffusion/diffusion_kv/test_block_tables.py b/tests/diffusion/diffusion_kv/test_block_tables.py index 6f88f2b4e..f2002f1c6 100644 --- a/tests/diffusion/diffusion_kv/test_block_tables.py +++ b/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( diff --git a/tests/diffusion/diffusion_kv/test_config.py b/tests/diffusion/diffusion_kv/test_config.py index a3b25dcd5..61b1b3041 100644 --- a/tests/diffusion/diffusion_kv/test_config.py +++ b/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", diff --git a/tests/diffusion/diffusion_kv/test_initialization.py b/tests/diffusion/diffusion_kv/test_initialization.py index d91d6fdd1..a82dc01ff 100644 --- a/tests/diffusion/diffusion_kv/test_initialization.py +++ b/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 diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py index 759fe7e46..1f242420d 100644 --- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py +++ b/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) diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py index 51ab6567d..ef5d75f4d 100644 --- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py +++ b/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, diff --git a/tests/diffusion/diffusion_kv/test_spec_discovery.py b/tests/diffusion/diffusion_kv/test_spec_discovery.py index 7c6d94a0e..31faf4d08 100644 --- a/tests/diffusion/diffusion_kv/test_spec_discovery.py +++ b/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)) diff --git a/tests/diffusion/diffusion_kv/test_worker_contract.py b/tests/diffusion/diffusion_kv/test_worker_contract.py index e0f428fb6..5f2caab06 100644 --- a/tests/diffusion/diffusion_kv/test_worker_contract.py +++ b/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() diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py index dfb383e87..b62fe950e 100644 --- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py +++ b/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: {}) diff --git a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py index 98026564c..65fc0e67b 100644 --- a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py +++ b/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] diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index d2a9f1165..814450d92 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/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() diff --git a/tests/platforms/npu/test_diffusion_platform.py b/tests/platforms/npu/test_diffusion_platform.py new file mode 100644 index 000000000..10e43e034 --- /dev/null +++ b/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 diff --git a/tests/platforms/test_record_device_event.py b/tests/platforms/test_record_device_event.py index 6a9e0ce06..9b81f9cd4 100644 --- a/tests/platforms/test_record_device_event.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index dd344e85d..a376216fa 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/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. diff --git a/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py b/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py index caa942987..a8676c34a 100644 --- a/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 2ab4c6f91..7dca3c1e6 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/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: diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 582346828..4c9c37c35 100644 --- a/vllm_omni/diffusion/data.py +++ b/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: diff --git a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py index 479492bac..c5ab3fb16 100644 --- a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py +++ b/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}") diff --git a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py index 5778af810..6a7862669 100644 --- a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py +++ b/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 diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index b8fddcfe8..8924e175e 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/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: diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 2cc42df71..a1ad2c986 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/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"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 4b1da8911..ce2d1d7e2 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/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): diff --git a/vllm_omni/diffusion/vllm_config.py b/vllm_omni/diffusion/vllm_config.py index 85e288422..ac591d687 100644 --- a/vllm_omni/diffusion/vllm_config.py +++ b/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 diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 80c3a1bc4..ea0dff90a 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 28a263e1f..5f8fa3e82 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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), diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 30421fa58..7ab4dd084 100644 --- a/vllm_omni/platforms/interface.py +++ b/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, diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 8316b067f..cf65da414 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/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: