diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index cee4a8214..f45f62ec1 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1090,7 +1090,10 @@ def test_diffusion_config_preserves_existing_coercion_hooks(): def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch): + from vllm_omni.platforms import current_omni_platform + monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE") + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) cfg = omni_config_module._DiffusionConfigProjection.from_kwargs( diffusion_attention_backend="flash_attn", @@ -1100,6 +1103,7 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch): kv_cache_skip_layers=[2], static_lora_scale=0.25, diffusion_kv_mode="paged_scheduler", + diffusion_kv_max_rows_per_request=2, diffusers_load_kwargs=None, diffusers_call_kwargs=None, ) @@ -1116,7 +1120,10 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch): assert cfg.diffusers_call_kwargs == {} -def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path): +def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch): + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) deploy_path = tmp_path / "dreamzero_diffusion_aliases.yaml" deploy_path.write_text( "\n".join( @@ -1128,6 +1135,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar " diffusion_attention_backend: flash_attn", " fa_deterministic: true", " diffusion_kv_mode: paged_scheduler", + " diffusion_kv_max_rows_per_request: 2", ] ) ) @@ -1141,6 +1149,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn" assert stage.diffusion_config.fa_deterministic is True assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + assert stage.diffusion_config.diffusion_kv_max_rows_per_request == 2 def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path): @@ -1183,6 +1192,7 @@ def test_diffusion_config_field_classification_covers_current_fields(): "prompt_embed_cache_size", "diffusion_kv_cache_dtype", "diffusion_kv_mode", + "diffusion_kv_max_rows_per_request", } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS assert { "revision", diff --git a/tests/diffusion/attention/test_piecewise_attn.py b/tests/diffusion/attention/test_piecewise_attn.py index 05967e28f..d1c414390 100644 --- a/tests/diffusion/attention/test_piecewise_attn.py +++ b/tests/diffusion/attention/test_piecewise_attn.py @@ -23,8 +23,10 @@ import vllm_omni.diffusion.attention.backends.utils.piecewise_attn as piecewise_ from vllm_omni.diffusion.attention.backends.abstract import QueryRange from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( Segment, + build_paged_piecewise_plan, build_segments, piecewise_attn, + run_paged_piecewise_plan, ) pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -84,6 +86,72 @@ def test_build_segments_full_span_retains_global_kv_end(): assert segments == [Segment(q_start=4, q_end=6, kv_end=7, mode="full")] +def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints(): + full_attn_spans = [[(2, 5)], [(4, 7)]] + plan = build_paged_piecewise_plan( + full_attn_spans, + query_offsets=[0, 2], + query_lens=[6, 6], + seq_lens=[6, 8], + device=DEVICE, + ) + query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4) + key = query + 100 + value = query + 200 + calls = [] + + def run_segment(segment_query, segment_key, segment_value, segment_batch): + calls.append( + ( + [segment.mode for segment in segment_batch.row_segments], + [segment.kv_end for segment in segment_batch.row_segments], + segment_batch.query_indices.tolist(), + ) + ) + torch.testing.assert_close(segment_key, segment_query + 100) + torch.testing.assert_close(segment_value, segment_query + 200) + return segment_query + + output = run_paged_piecewise_plan( + query, + key, + value, + plan, + plan.segments, + run_segment, + ) + + torch.testing.assert_close(output, query) + assert calls == [ + (["causal", "causal"], [2, 4], [0, 1, 6, 7]), + (["full", "full"], [5, 7], [2, 3, 4, 8, 9, 10]), + (["causal", "causal"], [6, 8], [5, 11]), + ] + + +def test_paged_piecewise_runner_uses_runner_output_shape(): + 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 = run_paged_piecewise_plan( + query, + query, + query, + plan, + plan.segments, + lambda segment_query, *_args: segment_query[..., :3], + ) + + assert output.shape == (6, 2, 3) + torch.testing.assert_close(output, query[..., :3]) + + @pytest.mark.parametrize("global_spans", SPAN_CASES) @pytest.mark.parametrize("q_range", Q_RANGE_CASES) @pytest.mark.parametrize("batch_size", BATCH_CASES) diff --git a/tests/diffusion/diffusion_kv/test_block_tables.py b/tests/diffusion/diffusion_kv/test_block_tables.py new file mode 100644 index 000000000..6f88f2b4e --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_block_tables.py @@ -0,0 +1,772 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import math +from contextlib import nullcontext +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from vllm.v1.kv_cache_interface import FullAttentionSpec +from vllm.v1.worker import block_table as native_block_table + +from vllm_omni.diffusion.diffusion_kv import model_runner_backend as model_runner_backend_module +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.metadata import ( + DiffusionKVContextMetadata, + DiffusionKVMetadata, + DiffusionKVSequenceMetadata, +) +from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend +from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput +from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _registration_backend(parallel_config: object): + model_config = SimpleNamespace(set_attention_geometry=Mock()) + vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(use_non_causal=False), + compilation_config=SimpleNamespace(static_forward_context={}), + model_config=model_config, + ) + backend = DiffusionKVModelRunnerBackend( + vllm_config=vllm_config, + od_config=SimpleNamespace(parallel_config=parallel_config), + device=torch.device("cpu"), + ) + return backend, vllm_config, model_config + + +def _parallel_config(**overrides): + values = dict(ulysses_degree=1, ring_degree=1, allgather_degree=1, ulysses_mode="strict") + values.update(overrides) + return SimpleNamespace(**values) + + +def test_cache_layer_registration_installs_adapters_with_ulysses_local_geometry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend, vllm_config, model_config = _registration_backend(_parallel_config(ulysses_degree=2)) + layer = SimpleNamespace( + num_heads=8, + softmax_scale=0.125, + skip_sequence_parallel=False, + attn_backend=SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN"), + ) + unrelated = object() + vllm_config.compilation_config.static_forward_context.update({"layer-0": layer, "unrelated": unrelated}) + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=8, + dtype=torch.float16, + non_causal=True, + ) + constructor_calls = [] + + class FakeLayerAdapter: + def __init__(self, *, layer_name, layer, spec, ulysses_degree, **_kwargs) -> None: + constructor_calls.append((layer_name, layer, spec, ulysses_degree)) + self.num_heads = layer.num_heads // ulysses_degree + self.num_kv_heads = spec.num_kv_heads // ulysses_degree + self.head_size = spec.head_size + self.spec = replace(spec, num_kv_heads=self.num_kv_heads) + + monkeypatch.setattr(model_runner_backend_module, "DiffusionPagedAttentionLayerAdapter", FakeLayerAdapter) + monkeypatch.setattr(model_runner_backend_module, "set_current_vllm_config", lambda _config: nullcontext()) + + specs = backend.register_kv_cache_layers({"layer-0": (layer, spec)}) + + installed = vllm_config.compilation_config.static_forward_context["layer-0"] + assert isinstance(installed, FakeLayerAdapter) + assert vllm_config.compilation_config.static_forward_context["unrelated"] is unrelated + assert specs["layer-0"].num_kv_heads == 2 + assert constructor_calls == [("layer-0", layer, spec, 2)] + model_config.set_attention_geometry.assert_called_once_with( + num_heads=4, + num_kv_heads=2, + head_size=8, + ) + assert vllm_config.attention_config.use_non_causal is True + + +def test_cache_layer_registration_rejects_backend_without_paged_support() -> None: + backend, _, _ = _registration_backend(_parallel_config()) + layer = SimpleNamespace( + attn_backend=SimpleNamespace(supports_paged_kv=False, get_name=lambda: "SDPA"), + ) + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=8, + dtype=torch.float16, + ) + + with pytest.raises(NotImplementedError, match="layer-0.*paged support.*SDPA"): + backend.register_kv_cache_layers({"layer-0": (layer, spec)}) + + +def test_cache_layer_registration_resolves_platform_paged_hooks(monkeypatch: pytest.MonkeyPatch) -> None: + backend, _, _ = _registration_backend(_parallel_config()) + + def unsupported_block_tables(): + raise NotImplementedError("platform has no paged BlockTables") + + monkeypatch.setattr( + model_runner_backend_module.current_omni_platform, + "get_diffusion_kv_block_tables_cls", + unsupported_block_tables, + ) + + with pytest.raises(NotImplementedError, match="no paged BlockTables"): + backend.register_kv_cache_layers({"layer-0": (object(), Mock(spec=FullAttentionSpec))}) + + +@pytest.mark.parametrize( + ("parallel_config", "message"), + [ + (_parallel_config(ring_degree=2), "Ring"), + (_parallel_config(allgather_degree=2), "AllGather-KV"), + (_parallel_config(ulysses_degree=2, ulysses_mode="advanced_uaa"), "strict Ulysses"), + ], +) +def test_cache_layer_registration_rejects_unsupported_sp_modes(parallel_config, message: str) -> None: + backend, _, _ = _registration_backend(parallel_config) + + with pytest.raises(NotImplementedError, match=message): + backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=False)) + + +def test_cache_layer_registration_ignores_sp_for_opted_out_layer() -> None: + backend, _, _ = _registration_backend( + _parallel_config(ulysses_degree=2, ring_degree=2, ulysses_mode="advanced_uaa") + ) + + assert backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=True)) == 1 + + +class _FakeSpec: + def __init__(self, block_size: int) -> None: + self.block_size = block_size + self.capacity_calls: list[tuple[object, int]] = [] + + def max_num_blocks_per_req(self, vllm_config: object, max_len: int) -> int: + self.capacity_calls.append((vllm_config, max_len)) + return math.ceil(max_len / self.block_size) + + +class _FakeStagedTable: + def __init__(self) -> None: + self.staged: list[tuple[int, tuple[int, ...]]] = [] + self.clear_calls = 0 + + def clear_staged_writes(self) -> None: + self.staged.clear() + self.clear_calls += 1 + + +class _FakeNumBlocks: + def __init__(self, num_groups: int, max_num_reqs: int) -> None: + self.np = np.zeros((num_groups, max_num_reqs), dtype=np.int32) + self.copy_calls = 0 + + def copy_to_uva(self) -> None: + self.copy_calls += 1 + + +class _FakeBlockTables: + instances: list[_FakeBlockTables] = [] + + def __init__(self, **kwargs: object) -> None: + self.kwargs = kwargs + self.block_sizes = list(kwargs["block_sizes"]) + self.max_num_reqs = int(kwargs["max_num_reqs"]) + self.max_num_batched_tokens = int(kwargs["max_num_batched_tokens"]) + self.cp_size = int(kwargs["cp_size"]) + self.cp_rank = int(kwargs["cp_rank"]) + self.cp_interleave = int(kwargs["cp_interleave"]) + self.blocks_per_kv_block = [1] * len(self.block_sizes) + self.block_tables = [_FakeStagedTable() for _ in self.block_sizes] + self.num_blocks = _FakeNumBlocks(len(self.block_sizes), self.max_num_reqs) + self.append_calls: list[tuple[int, tuple[list[int], ...], bool]] = [] + self.apply_calls = 0 + self.layout_refreshes = 0 + self.fail_on_append_call: int | None = None + self.fail_apply = False + self.committed: list[dict[int, tuple[int, ...]]] = [dict() for _ in self.block_sizes] + self.__class__.instances.append(self) + + def append_block_ids( + self, + req_index: int, + new_block_ids: tuple[list[int], ...], + overwrite: bool, + ) -> None: + call_number = len(self.append_calls) + 1 + if self.fail_on_append_call == call_number: + raise RuntimeError("fake append failure") + copied_ids = tuple(list(group_ids) for group_ids in new_block_ids) + self.append_calls.append((req_index, copied_ids, overwrite)) + for group_index, group_ids in enumerate(copied_ids): + self.num_blocks.np[group_index, req_index] = len(group_ids) + if group_ids: + self.block_tables[group_index].staged.append((req_index, tuple(group_ids))) + + def apply_staged_writes(self) -> None: + self.apply_calls += 1 + if self.fail_apply: + raise RuntimeError("fake apply failure") + for group_index, group_table in enumerate(self.block_tables): + for row, group_ids in group_table.staged: + self.committed[group_index][row] = group_ids + group_table.staged.clear() + self.num_blocks.copy_to_uva() + + def init_block_table_layout_tensors(self) -> None: + self.layout_refreshes += 1 + + +def _metadata( + request_id: str = "req-0", + generation: int = 1, + *, + seq_len: int = 9, + sequence_block_ids: tuple[list[int], ...] = ([1, 2, 3], [4, 5]), + contexts: tuple[DiffusionKVContextMetadata, ...] | None = None, +) -> DiffusionKVMetadata: + if contexts is None: + contexts = ( + DiffusionKVContextMetadata( + context_id="text", + cache_role="cross_attention", + num_tokens=5, + block_ids=([6, 7], [8]), + ), + ) + return DiffusionKVMetadata( + request_id=request_id, + allocation_generation=generation, + sequences=( + DiffusionKVSequenceMetadata( + sequence_id=0, + prefix_len=4, + target_len=5, + seq_len=seq_len, + block_ids=sequence_block_ids, + context_ids=tuple(context.context_id for context in contexts), + ), + ), + contexts=contexts, + ) + + +def _install_native_modules(monkeypatch: pytest.MonkeyPatch, events: list[tuple]) -> None: + _FakeBlockTables.instances.clear() + + def get_block_table_width(max_num_blocks, block_size): + events.append(("width", max_num_blocks, block_size)) + return max_num_blocks + 1 + + def init_backend(kv_cache_config, vllm_config, device): + events.append(("backend", kv_cache_config, vllm_config, device)) + group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None)) + return [[group], [group]], "cg-support", [4, 8] + + def init_cache( + runner_kv_caches, + forward_context, + kv_cache_config, + attn_groups, + device, + cache_dtype, + kernel_block_sizes, + vllm_config, + ): + events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes)) + runner_kv_caches.extend(["cache-0", "cache-1"]) + return {"layer-0": "tensor-0", "layer-1": "tensor-1"} + + monkeypatch.setattr(model_runner_backend_module, "init_attn_backend", init_backend) + monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", init_cache) + monkeypatch.setattr( + model_runner_backend_module.current_omni_platform, + "get_diffusion_kv_block_tables_cls", + lambda: _FakeBlockTables, + ) + monkeypatch.setattr(native_block_table, "get_block_table_width", get_block_table_width, raising=False) + + +def _runner( + monkeypatch: pytest.MonkeyPatch, + *, + max_rows_per_request: int = 3, + max_num_seqs: int = 2, + max_len: int = 16, + max_num_batched_tokens: int = 32, + cp_size: int = 1, + cp_rank: int = 0, + cp_interleave: int = 1, +) -> tuple[DiffusionKVModelRunnerBackend, _FakeBlockTables, list[tuple]]: + events: list[tuple] = [] + _install_native_modules(monkeypatch, events) + od_config = SimpleNamespace( + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + diffusion_kv_max_rows_per_request=max_rows_per_request, + max_num_seqs=max_num_seqs, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=max_len), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=max_num_batched_tokens, + max_num_seqs=max_num_seqs, + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=cp_size, + cp_kv_cache_interleave_size=cp_interleave, + ), + compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}), + cache_config=SimpleNamespace(cache_dtype="auto"), + ) + monkeypatch.setattr( + model_runner_backend_module, + "get_dcp_group", + lambda: SimpleNamespace(rank_in_group=cp_rank), + ) + backend = DiffusionKVModelRunnerBackend( + vllm_config=vllm_config, + od_config=od_config, + device="cuda:1", + ) + backend._kv_cache_layer_adapters = { + "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None), + "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None), + } + + config = SimpleNamespace( + num_blocks=32, + kv_cache_groups=[ + SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)), + SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)), + ], + ) + backend.initialize_kv_cache(config) + return backend, _FakeBlockTables.instances[-1], events + + +def test_initialize_builds_native_block_tables_from_rank_local_config(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, events = _runner(monkeypatch) + + assert events[:2] == [("width", 4, 4), ("width", 2, 8)] + assert [event[0] for event in events[2:]] == ["backend", "cache"] + assert block_tables.kwargs == { + "block_sizes": [4, 8], + "max_num_reqs": 6, + "max_num_batched_tokens": 32, + "max_num_blocks_per_group": [5, 3], + "device": "cuda:1", + "kernel_block_sizes": [4, 8], + "cp_size": 1, + "cp_rank": 0, + "cp_interleave": 1, + } + assert runner.block_tables is block_tables + assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0] + assert runner.kv_caches == ["cache-0", "cache-1"] + assert runner.vllm_config.scheduler_config.max_num_seqs == 6 + + +def test_initialize_propagates_native_context_parallel_layout(monkeypatch: pytest.MonkeyPatch) -> None: + _, block_tables, _ = _runner( + monkeypatch, + cp_size=2, + cp_rank=1, + cp_interleave=4, + ) + + assert block_tables.kwargs["cp_size"] == 2 + assert block_tables.kwargs["cp_rank"] == 1 + assert block_tables.kwargs["cp_interleave"] == 4 + + +def test_initialize_failure_does_not_publish_partial_cache_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[tuple] = [] + _install_native_modules(monkeypatch, events) + successful_init_cache = model_runner_backend_module.init_kv_cache + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=16), + scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=1), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + cp_kv_cache_interleave_size=1, + ), + compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}), + cache_config=SimpleNamespace(cache_dtype="auto"), + ) + backend = DiffusionKVModelRunnerBackend( + vllm_config=vllm_config, + od_config=SimpleNamespace( + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + diffusion_kv_max_rows_per_request=2, + ), + device="cuda:1", + ) + backend._kv_cache_layer_adapters = { + "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-0"), + "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-1"), + } + config = SimpleNamespace( + num_blocks=32, + kv_cache_groups=[ + SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)), + SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)), + ], + ) + + def fail_init_cache(runner_kv_caches, *_args, **_kwargs): + runner_kv_caches.append("partial-cache") + backend._kv_cache_layer_adapters["layer-0"].kv_cache = "partial-binding" + raise RuntimeError("cache binding failed") + + monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", fail_init_cache) + with pytest.raises(RuntimeError, match="cache binding failed"): + backend.initialize_kv_cache(config) + + assert backend.kv_cache_config is None + assert backend.kv_caches == [] + assert backend.block_tables is None + assert backend.paged_attention_adapter is None + assert backend._kv_cache_layer_adapters["layer-0"].kv_cache == "placeholder-0" + assert backend.vllm_config.scheduler_config.max_num_seqs == 1 + + monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", successful_init_cache) + backend.initialize_kv_cache(config) + assert backend.kv_cache_config is not None + assert backend.kv_caches == ["cache-0", "cache-1"] + assert backend.vllm_config.scheduler_config.max_num_seqs == 2 + + +def test_initialize_rejects_config_for_different_registered_layers() -> None: + backend = DiffusionKVModelRunnerBackend( + vllm_config=SimpleNamespace(), + od_config=SimpleNamespace(), + device="cuda:1", + ) + backend._kv_cache_layer_adapters = { + "expected-layer": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None) + } + config = SimpleNamespace( + kv_cache_groups=[SimpleNamespace(layer_names=["other-layer"])], + ) + + with pytest.raises(ValueError, match="layer mismatch"): + backend.initialize_kv_cache(config) + + +def test_valid_sequence_and_context_install_into_native_rows(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + + assert runner.install_diffusion_kv_metadata(_metadata()) is True + + assert block_tables.append_calls == [ + (0, ([1, 2, 3], [4, 5]), True), + (1, ([6, 7], [8]), True), + ] + assert block_tables.apply_calls == 1 + assert runner.get_diffusion_kv_row("req-0", 0) == 0 + 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) + + +def test_block_table_mutations_invalidate_prepared_attention_batches( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner, _, _ = _runner(monkeypatch) + assert runner.paged_attention_adapter is not None + runner.paged_attention_adapter.invalidate_prepared_batches = Mock() + + runner.install_diffusion_kv_metadata(_metadata()) + runner.paged_attention_adapter.invalidate_prepared_batches.assert_called_once_with() + + runner.remove_diffusion_kv_requests(["req-0"]) + assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 2 + + runner.refresh_block_table_layout() + assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 3 + + +def test_cache_layer_registration_is_frozen_after_physical_initialization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner, _, _ = _runner(monkeypatch) + + with pytest.raises(RuntimeError, match="cannot be changed"): + runner.register_kv_cache_layers({}) + + +def test_request_scoped_context_uses_one_row_across_sequences(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + first_sequence = _metadata().sequences[0] + metadata = replace( + _metadata(), + sequences=( + first_sequence, + replace( + first_sequence, + sequence_id=1, + block_ids=([9, 10, 11], [12, 13]), + ), + ), + ) + + assert runner.install_diffusion_kv_metadata(metadata) is True + + assert block_tables.append_calls == [ + (0, ([1, 2, 3], [4, 5]), True), + (1, ([9, 10, 11], [12, 13]), True), + (2, ([6, 7], [8]), True), + ] + assert runner.get_diffusion_kv_row("req-0", 0, "text") == 2 + assert runner.get_diffusion_kv_row("req-0", 1, "text") == 2 + + +def test_row_block_count_uses_native_spec_semantics(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + spec = runner.kv_cache_config.kv_cache_groups[0].kv_cache_spec + spec.max_num_blocks_per_req = lambda _config, token_len: math.ceil(token_len / 8) + + metadata = _metadata( + sequence_block_ids=([1, 2], [4, 5]), + contexts=(), + ) + + assert runner.install_diffusion_kv_metadata(metadata) is True + assert block_tables.append_calls == [(0, ([1, 2], [4, 5]), True)] + + +def test_repeated_native_block_ids_are_allowed(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + metadata = _metadata( + sequence_block_ids=([1, 1, 3], [4, 4]), + contexts=(), + ) + + assert runner.install_diffusion_kv_metadata(metadata) is True + assert block_tables.append_calls == [(0, ([1, 1, 3], [4, 4]), True)] + + +def test_generation_is_idempotent_stale_safe_and_conflict_safe(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + metadata = _metadata() + runner.install_diffusion_kv_metadata(metadata) + + assert runner.install_diffusion_kv_metadata(metadata) is False + assert runner.install_diffusion_kv_metadata(replace(metadata, allocation_generation=2)) is False + assert block_tables.apply_calls == 1 + + with pytest.raises(ValueError, match="Stale Diffusion KV allocation generation"): + runner.install_diffusion_kv_metadata(metadata) + with pytest.raises(ValueError, match="Conflicting Diffusion KV allocation snapshot"): + runner.install_diffusion_kv_metadata(_metadata(generation=3, sequence_block_ids=([9, 2, 3], [4, 5]))) + assert block_tables.apply_calls == 1 + + +@pytest.mark.parametrize( + ("metadata", "message"), + [ + (_metadata(sequence_block_ids=([1, 2, 3],)), "block groups"), + (_metadata(sequence_block_ids=([1, 2], [4, 5])), "expected 3"), + (_metadata(sequence_block_ids=([1, 2, 32], [4, 5])), "outside"), + ( + _metadata( + seq_len=17, + sequence_block_ids=([1, 2, 3, 4, 5], [6, 7, 8]), + contexts=(), + ), + "exceeds row capacity", + ), + ], +) +def test_invalid_group_count_block_count_range_and_row_capacity_are_atomic( + monkeypatch: pytest.MonkeyPatch, + metadata: DiffusionKVMetadata, + message: str, +) -> None: + runner, block_tables, _ = _runner(monkeypatch) + + with pytest.raises(ValueError, match=message): + runner.install_diffusion_kv_metadata(metadata) + + assert block_tables.append_calls == [] + assert block_tables.apply_calls == 0 + assert runner._diffusion_kv_identity_to_row == {} + assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0] + + +def test_invalid_later_context_does_not_stage_earlier_valid_row(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + bad_context = DiffusionKVContextMetadata( + context_id="text", + cache_role="cross_attention", + num_tokens=5, + block_ids=([6, 32], [8]), + ) + + with pytest.raises(ValueError, match="outside"): + runner.install_diffusion_kv_metadata(_metadata(contexts=(bad_context,))) + + assert block_tables.append_calls == [] + assert all(not group_table.staged for group_table in block_tables.block_tables) + + +def test_native_append_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + block_tables.fail_on_append_call = 2 + + with pytest.raises(RuntimeError, match="fake append failure"): + runner.install_diffusion_kv_metadata(_metadata()) + + assert np.count_nonzero(block_tables.num_blocks.np) == 0 + assert all(not group_table.staged for group_table in block_tables.block_tables) + assert runner._diffusion_kv_identity_to_row == {} + assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0] + + +def test_native_apply_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + block_tables.fail_apply = True + + with pytest.raises(RuntimeError, match="fake apply failure"): + runner.install_diffusion_kv_metadata(_metadata()) + + assert block_tables.apply_calls == 1 + assert np.count_nonzero(block_tables.num_blocks.np) == 0 + assert all(not group_table.staged for group_table in block_tables.block_tables) + assert all(group_table.clear_calls == 1 for group_table in block_tables.block_tables) + assert runner._diffusion_kv_identity_to_row == {} + assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0] + + +def test_adapter_and_global_row_capacity_are_enforced(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1) + second_context = DiffusionKVContextMetadata( + context_id="image", + cache_role="cross_attention", + num_tokens=1, + block_ids=([9], [10]), + ) + + with pytest.raises(ValueError, match="adapter limit is 2"): + runner.install_diffusion_kv_metadata(_metadata(contexts=(_metadata().contexts[0], second_context))) + assert block_tables.append_calls == [] + + runner.install_diffusion_kv_metadata(_metadata()) + append_count = len(block_tables.append_calls) + with pytest.raises(ValueError, match="only 0 rows are free"): + runner.install_diffusion_kv_metadata(_metadata("req-1")) + assert len(block_tables.append_calls) == append_count + + +def test_empty_and_duplicate_sequence_metadata_are_rejected(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + empty = DiffusionKVMetadata(request_id="req-0", allocation_generation=0, sequences=()) + duplicate_sequence = _metadata().sequences[0] + duplicate = DiffusionKVMetadata( + request_id="req-0", + allocation_generation=0, + sequences=(duplicate_sequence, duplicate_sequence), + ) + + with pytest.raises(ValueError, match="at least one sequence"): + runner.install_diffusion_kv_metadata(empty) + with pytest.raises(ValueError, match="Duplicate diffusion KV row identity"): + runner.install_diffusion_kv_metadata(duplicate) + assert block_tables.append_calls == [] + + +@pytest.mark.parametrize( + ("metadata", "message"), + [ + (_metadata(generation=-1), "non-negative integer"), + (_metadata(request_id=""), "non-empty string"), + (_metadata(sequence_block_ids=([1, 2, True], [4, 5])), "non-integer block ID"), + ], +) +def test_request_generation_and_block_id_types_are_validated( + monkeypatch: pytest.MonkeyPatch, + metadata: DiffusionKVMetadata, + message: str, +) -> None: + runner, block_tables, _ = _runner(monkeypatch) + + with pytest.raises(ValueError, match=message): + runner.install_diffusion_kv_metadata(metadata) + + assert block_tables.append_calls == [] + assert runner._diffusion_kv_request_states == {} + + +def test_cleanup_is_idempotent_and_reuses_native_rows(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1) + runner.install_diffusion_kv_metadata(_metadata()) + + assert runner.remove_diffusion_kv_requests(["missing"]) == 0 + assert runner.remove_diffusion_kv_requests(["req-0", "req-0"]) == 2 + assert block_tables.append_calls[-2:] == [(0, ([], []), True), (1, ([], []), True)] + assert block_tables.num_blocks.np[:, :2].tolist() == [[0, 0], [0, 0]] + assert runner.remove_diffusion_kv_requests(["req-0"]) == 0 + with pytest.raises(KeyError, match="No native diffusion KV row"): + runner.get_diffusion_kv_row("req-0", 0) + + runner.install_diffusion_kv_metadata(_metadata("req-1")) + assert runner.get_diffusion_kv_row("req-1", 0) == 0 + assert runner.get_diffusion_kv_row("req-1", 0, "text") == 1 + + +def test_cached_step_does_not_reinstall_metadata() -> None: + runner = object.__new__(DiffusionModelRunner) + runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER) + runner.pipeline = object() + runner.install_diffusion_kv_metadata = Mock() + runner._supports_step_mode = Mock(return_value=False) + scheduler_output = DiffusionSchedulerOutput( + step_id=1, + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData(request_ids=["req-0"]), + finished_req_ids=set(), + num_running_reqs=1, + num_waiting_reqs=0, + ) + + with pytest.raises(ValueError, match="does not support step execution"): + runner.execute_stepwise(scheduler_output) + runner.install_diffusion_kv_metadata.assert_not_called() + + +def test_dense_cleanup_and_wake_are_noops_but_paged_wake_refreshes(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + runner.refresh_block_table_layout() + assert block_tables.layout_refreshes == 1 + + dense_backend = DiffusionKVModelRunnerBackend( + vllm_config=SimpleNamespace(), + od_config=SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY), + device="cuda:1", + ) + dense_backend.block_tables = Mock() + assert dense_backend.remove_diffusion_kv_requests(["req-0"]) == 0 + dense_backend.refresh_block_table_layout() + dense_backend.block_tables.init_block_table_layout_tensors.assert_not_called() + with pytest.raises(ValueError, match="Dense diffusion execution"): + dense_backend.install_diffusion_kv_metadata(_metadata()) diff --git a/tests/diffusion/diffusion_kv/test_config.py b/tests/diffusion/diffusion_kv/test_config.py index f03120784..a3b25dcd5 100644 --- a/tests/diffusion/diffusion_kv/test_config.py +++ b/tests/diffusion/diffusion_kv/test_config.py @@ -3,6 +3,7 @@ import pytest +from vllm_omni.config.omni_config import _DiffusionConfigProjection from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode @@ -39,6 +40,8 @@ def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None: with pytest.raises(ValueError, match="does not support imported AR KV"): OmniDiffusionConfig.from_kwargs( diffusion_kv_mode="paged_scheduler", + diffusion_kv_max_rows_per_request=1, + max_num_batched_tokens=1, omni_kv_config={"need_recv_cache": True}, ) @@ -47,6 +50,33 @@ def test_paged_scheduler_does_not_depend_on_model_registry() -> None: config = OmniDiffusionConfig.from_kwargs( model_class_name="FutureDiffusionModel", diffusion_kv_mode="paged_scheduler", + diffusion_kv_max_rows_per_request=1, + max_num_batched_tokens=1, ) assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + + +@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection]) +def test_paged_scheduler_mode_is_platform_agnostic(config_cls) -> None: + config = config_cls.from_kwargs( + diffusion_kv_mode="paged_scheduler", + diffusion_kv_max_rows_per_request=1, + ) + + assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + + +def test_paged_scheduler_requires_a_worker_row_limit() -> None: + with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"): + OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_scheduler") + with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"): + _DiffusionConfigProjection.from_kwargs(diffusion_kv_mode="paged_scheduler") + + +@pytest.mark.parametrize("invalid_limit", [0, -1, True, 1.5]) +def test_worker_row_limit_must_be_a_positive_integer(invalid_limit: object) -> None: + with pytest.raises((TypeError, ValueError)): + OmniDiffusionConfig.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit) + with pytest.raises((TypeError, ValueError)): + _DiffusionConfigProjection.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit) diff --git a/tests/diffusion/diffusion_kv/test_initialization.py b/tests/diffusion/diffusion_kv/test_initialization.py index 126b25d18..d91d6fdd1 100644 --- a/tests/diffusion/diffusion_kv/test_initialization.py +++ b/tests/diffusion/diffusion_kv/test_initialization.py @@ -22,6 +22,7 @@ class _Executor: self.specs = [{"layer0": object()}, {"layer0": object()}] self.memory = [1024, 1024] self.configs = None + self.resolved_max_model_len = None self.profile_requests = None def get_kv_cache_specs(self): @@ -31,8 +32,9 @@ class _Executor: self.profile_requests = profile_requests return self.memory - def set_kv_cache_configs(self, configs) -> None: + def set_kv_cache_configs(self, configs, resolved_max_model_len) -> None: self.configs = configs + self.resolved_max_model_len = resolved_max_model_len def _od_config(**overrides): @@ -55,6 +57,7 @@ def _od_config(**overrides): ), gpu_memory_utilization=0.9, kv_cache_memory_bytes=None, + max_num_seqs=1, max_num_batched_tokens=64, num_gpus=1, ) @@ -64,9 +67,9 @@ def _od_config(**overrides): def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None: executor = _Executor() - od_config = _od_config(num_gpus=2) + od_config = _od_config(num_gpus=2, max_model_len=-1) vllm_config = SimpleNamespace( - model_config=SimpleNamespace(max_model_len=4096), + model_config=SimpleNamespace(max_model_len=256), max_in_flight_tokens=128, ) worker_configs = [object(), object()] @@ -91,6 +94,7 @@ def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None: build.assert_called_once_with(vllm_config, executor.specs, executor.memory) assert result == (scheduler_kv_cache_config, 16, 16, vllm_config) assert executor.configs == worker_configs + assert executor.resolved_max_model_len == vllm_config.model_config.max_model_len assert executor.profile_requests is profile_requests diff --git a/tests/diffusion/diffusion_kv/test_native_block_tables_gpu.py b/tests/diffusion/diffusion_kv/test_native_block_tables_gpu.py new file mode 100644 index 000000000..4848fd49a --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_native_block_tables_gpu.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from importlib import import_module, util + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] + + +def test_native_block_tables_slot_mapping_handles_non_aligned_tail() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the native BlockTables slot-mapping smoke test") + module_name = "vllm.v1.worker.gpu.block_table" + try: + module_spec = util.find_spec(module_name) + except ModuleNotFoundError as exc: + if exc.name is None or not module_name.startswith(f"{exc.name}."): + raise + pytest.skip("installed vLLM does not provide GPU BlockTables") + if module_spec is None: + pytest.skip("installed vLLM does not provide GPU BlockTables") + block_table_module = import_module(module_name) + block_tables_cls = getattr(block_table_module, "BlockTables", None) + if block_tables_cls is None or not hasattr(block_tables_cls, "compute_slot_mappings"): + pytest.skip("installed vLLM does not provide the required BlockTables API") + + device = torch.device("cuda", torch.accelerator.current_device_index()) + block_tables = block_tables_cls( + block_sizes=[4], + max_num_reqs=1, + max_num_batched_tokens=16, + max_num_blocks_per_group=[4], + device=device, + kernel_block_sizes=[4], + ) + block_tables.append_block_ids(0, ([1, 2, 3],), overwrite=True) + block_tables.apply_staged_writes() + idx_mapping = torch.tensor([0], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 9], dtype=torch.int32, device=device) + positions = torch.arange(9, dtype=torch.int64, device=device) + + slot_mappings = block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + num_tokens_padded=9, + ) + torch.accelerator.synchronize(device) + + assert slot_mappings.cpu().tolist() == [[4, 5, 6, 7, 8, 9, 10, 11, 12]] diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py new file mode 100644 index 000000000..759fe7e46 --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py @@ -0,0 +1,1147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import math +from contextlib import nullcontext +from dataclasses import dataclass +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from torch import nn +from vllm.v1.kv_cache_interface import FullAttentionSpec + +from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl +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, + DiffusionPagedAttentionRow, + DiffusionPagedAttentionRowBinding, +) +from vllm_omni.diffusion.forward_context import ( + override_paged_kv_adapter, + set_forward_context, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _FakeBlockTables: + def __init__(self) -> None: + self.max_num_reqs = 4 + self.max_num_batched_tokens = 16 + self.cp_size = 1 + self.cp_rank = 0 + self.cp_interleave = 1 + self.block_sizes = [4] + self.blocks_per_kv_block = [1] + self.num_blocks = SimpleNamespace(np=np.full((1, self.max_num_reqs), 4, dtype=np.int32)) + self.gather_calls: list[tuple[torch.Tensor, int]] = [] + self.slot_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]] = [] + + def gather_block_tables( + self, + idx_mapping: torch.Tensor, + num_reqs_padded: int, + ) -> tuple[torch.Tensor, ...]: + self.gather_calls.append((idx_mapping.clone(), num_reqs_padded)) + return (torch.tensor([[3, 4], [7, 8]], dtype=torch.int32),) + + def compute_slot_mappings( + self, + idx_mapping: torch.Tensor, + query_start_loc: torch.Tensor, + positions: torch.Tensor, + num_tokens_padded: int, + ) -> torch.Tensor: + self.slot_calls.append( + ( + idx_mapping.clone(), + query_start_loc.clone(), + positions.clone(), + num_tokens_padded, + ) + ) + return positions[:num_tokens_padded].to(torch.int64).unsqueeze(0) + + +class _FakeSpec: + non_causal: bool + block_size = 4 + + def __init__(self, *, non_causal: bool) -> None: + self.non_causal = non_causal + + def max_num_blocks_per_req(self, _vllm_config, max_len: int) -> int: + return math.ceil(max_len / self.block_size) + + +class _FakeAttentionGroup: + def __init__(self, reorder_batch_threshold: int | None) -> None: + self.builder = SimpleNamespace(reorder_batch_threshold=reorder_batch_threshold) + + def get_metadata_builder(self, _index: int): + return self.builder + + +@dataclass(frozen=True, slots=True) +class _FakeNativeMetadata: + build_id: int + causal: bool + seq_lens: torch.Tensor + query_start_loc_cpu: torch.Tensor + positions: torch.Tensor + slot_mappings: torch.Tensor + + +class _FakeLayer: + def __init__(self, *, non_causal: bool, num_heads: int = 2, num_kv_heads: int = 2) -> None: + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_size = 4 + self.head_size_v = 4 + self.spec = _FakeSpec(non_causal=non_causal) + self.updates: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, object]] = [] + self.native_events: list[str] = [] + self.layer_name = "layer-0" + self.kv_cache = object() + self.attn_backend = SimpleNamespace(forward_includes_kv_cache_update=False) + self.impl = _FakeNativeImpl(self) + + +class _FakeNativeImpl: + def __init__(self, layer: _FakeLayer) -> None: + self.layer = layer + + def do_kv_cache_update( + self, + _layer, + key: torch.Tensor, + value: torch.Tensor, + _kv_cache, + slot_mapping: torch.Tensor, + ) -> None: + self.layer.native_events.append("update") + self.layer.updates.append((key, value, slot_mapping)) + + def forward( + self, + _layer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + _kv_cache, + metadata: object, + output: torch.Tensor, + ) -> torch.Tensor: + self.layer.native_events.append("forward") + self.layer.calls.append((query, key, value, metadata)) + return output.copy_(query.reshape_as(output)) + + +def _make_adapter( + monkeypatch: pytest.MonkeyPatch, + *, + non_causal: bool = True, + reorder_batch_threshold: int | None = None, + num_heads: int = 2, + num_kv_heads: int = 2, + capacity: int = 16, +) -> tuple[DiffusionPagedAttentionAdapter, _FakeBlockTables, _FakeLayer, list[tuple]]: + events: list[tuple] = [] + block_tables = _FakeBlockTables() + block_tables.max_num_batched_tokens = capacity + block_tables.num_blocks.np.fill(math.ceil(capacity / block_tables.block_sizes[0])) + layer = _FakeLayer(non_causal=non_causal, num_heads=num_heads, num_kv_heads=num_kv_heads) + kv_cache_config = SimpleNamespace( + kv_cache_groups=[SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=layer.spec)], + ) + row_map = { + ("req-0", 0, None): DiffusionPagedAttentionRowBinding(2, capacity), + ("req-1", 1, None): DiffusionPagedAttentionRowBinding(3, capacity), + ("req-0", None, "text"): DiffusionPagedAttentionRowBinding(1, capacity), + } + + monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext()) + + def build_metadata(**kwargs): + metadata = _FakeNativeMetadata( + build_id=len(events), + 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(), + ) + events.append(("build", kwargs, metadata)) + return {"layer-0": metadata} + + monkeypatch.setattr(adapter_module.current_omni_platform, "build_diffusion_kv_attn_metadata", build_metadata) + monkeypatch.setattr( + adapter_module, + "build_slot_mappings_by_layer", + lambda slot_mappings, _config: {"layer-0": slot_mappings[0]}, + ) + + config = SimpleNamespace( + name="vllm-config", + model_config=SimpleNamespace(dtype=torch.float32), + ) + adapter = DiffusionPagedAttentionAdapter( + vllm_config=config, + device=torch.device("cpu"), + kv_cache_config=kv_cache_config, + block_tables=block_tables, + attn_groups=[[_FakeAttentionGroup(reorder_batch_threshold)]], + layers={"layer-0": layer}, + resolve_row=lambda request_id, sequence_id, context_id: row_map[(request_id, sequence_id, context_id)], + ) + return adapter, block_tables, layer, events + + +def _run_omni_paged_backend(context): + backend = FlashAttentionImpl( + num_heads=context.layer.num_heads, + head_size=context.layer.head_size, + softmax_scale=0.5, + num_kv_heads=context.layer.num_kv_heads, + ) + return backend.forward_paged(context) + + +def test_prepare_batch_reuses_native_block_table_and_metadata_builders(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, block_tables, _, events = _make_adapter(monkeypatch) + rows = ( + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=3, + seq_len=8, + kv_start_pos=4, + ), + DiffusionPagedAttentionRow( + request_id="req-1", + sequence_id=1, + query_len=2, + seq_len=2, + ), + ) + + batch = adapter.prepare_batch(rows) + + assert batch.num_tokens == 5 + assert batch.row_indices.tolist() == [2, 3] + assert batch.query_start_loc.tolist() == [0, 3, 5] + assert batch.seq_lens.tolist() == [8, 2] + assert batch.positions.tolist() == [4, 5, 6, 0, 1] + assert block_tables.gather_calls[0][1] == 2 + assert block_tables.gather_calls[0][0].tolist() == [2, 3] + assert block_tables.slot_calls[0][3] == 5 + build_kwargs = events[0][1] + assert build_kwargs["causal"] is False + assert build_kwargs["max_query_len"] == 3 + assert build_kwargs["max_seq_len"] == 8 + assert build_kwargs["positions"] is batch.positions + assert batch.attn_metadata["layer-0"] is events[0][2] + assert events[0][2].build_id == 0 + assert batch.slot_mappings_by_layer["layer-0"].tolist() == [4, 5, 6, 0, 1] + + +def test_omni_paged_backend_consumes_context_and_restores_diffusion_shape( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, layer, events = _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_like(query) + value = torch.randn_like(query) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", query, key, value) + output = _run_omni_paged_backend(context) + + assert output.shape == query.shape + assert torch.equal(output, query) + assert layer.calls[0][0].shape == (5, 2, 4) + assert layer.native_events == ["update", "forward"] + assert layer.calls[0][3] is events[0][2] + + +def test_paged_adapter_accepts_native_gqa_shapes(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch, num_heads=32, num_kv_heads=8) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)] + ) + query = torch.randn(1, 3, 32, 4) + key = torch.randn(1, 3, 8, 4) + value = torch.randn_like(key) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", query, key, value) + + assert layer.num_heads == 32 + assert layer.num_kv_heads == 8 + assert context.query.shape == (3, 32, 4) + assert context.key_write.shape == (3, 8, 4) + assert context.value_write.shape == (3, 8, 4) + + +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 + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)] + ) + qkv = torch.randn(1, 3, 2, 4) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + output = _run_omni_paged_backend(context) + + assert torch.equal(output, qkv) + assert layer.native_events == ["forward"] + assert layer.updates == [] + + +def test_omni_piecewise_backend_defers_backend_owned_cache_update(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, events = _make_adapter(monkeypatch) + layer.attn_backend.forward_includes_kv_cache_update = True + 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, + ) + output = _run_omni_paged_backend(context) + + assert torch.equal(output, qkv) + assert layer.native_events == ["forward", "forward", "forward"] + assert layer.updates == [] + assert [call[0].shape[0] for call in layer.calls] == [1, 2, 1] + assert [call[1].shape[0] for call in layer.calls] == [1, 2, 1] + assert [call[2].shape[0] for call in layer.calls] == [1, 2, 1] + assert all(call[3] is event[2] for call, event in zip(layer.calls, events[1:], strict=True)) + segment_metadata = [call[3] for call in layer.calls] + assert [metadata.causal for metadata in segment_metadata] == [True, False, True] + assert [metadata.seq_lens.tolist() for metadata in segment_metadata] == [[1], [3], [4]] + 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]]] + + +def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, events = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=6, + seq_len=9, + kv_start_pos=3, + ), + DiffusionPagedAttentionRow( + request_id="req-1", + sequence_id=1, + query_len=6, + seq_len=9, + kv_start_pos=3, + ), + ] + ) + query = torch.arange(2 * 6 * 2 * 4, dtype=torch.float32).reshape(2, 6, 2, 4) + key = torch.arange(2 * 6 * 2 * 4, dtype=torch.float32).reshape(2, 6, 2, 4) + 100 + value = key + 100 + attn_mask = torch.ones(9, 9, dtype=torch.bool).tril().repeat(2, 1, 1) + attn_mask[:, 5:8, 5:8] = True + metadata = SimpleNamespace( + attn_mask=attn_mask.unsqueeze(1), + full_attn_spans=[[(5, 8)], [(5, 8)]], + query_ranges=None, + extra={}, + ) + + with adapter.activate(batch): + context = adapter.prepare_layer_context( + "layer-0", + query, + key, + value, + omni_attn_metadata=metadata, + ) + # Every model layer sees the same spans, so the native segment metadata + # is built once per active batch and then reused. + adapter.prepare_layer_context( + "layer-0", + query, + key, + value, + omni_attn_metadata=metadata, + ) + output = _run_omni_paged_backend(context) + + assert torch.equal(output, query) + assert layer.native_events == ["update", "forward", "forward", "forward"] + assert len(layer.updates) == 1 + assert layer.updates[0][0].shape[0] == 12 + assert torch.equal(layer.updates[0][0], key.reshape(12, 2, 4)) + assert torch.equal(layer.updates[0][1], value.reshape(12, 2, 4)) + assert [call[0].shape[0] for call in layer.calls] == [4, 6, 2] + assert [call[1].shape[0] for call in layer.calls] == [4, 6, 2] + assert [call[2].shape[0] for call in layer.calls] == [4, 6, 2] + assert all(call[3] is event[2] for call, event in zip(layer.calls, events[1:], strict=True)) + assert [metadata.build_id for metadata in context.piecewise_native_metadata] == [1, 2, 3] + + # The first metadata build is the normal whole-query path. Piecewise calls + # then use causal [3, 5), full [5, 8), and causal [8, 9) segments. + assert len(events) == 4 + segment_builds = [event[1] for event in events[1:]] + assert [build["causal"] for build in segment_builds] == [True, False, True] + assert [build["seq_lens"].tolist() for build in segment_builds] == [ + [5, 5], + [8, 8], + [9, 9], + ] + assert [build["query_start_loc_cpu"].tolist() for build in segment_builds] == [ + [0, 2, 4], + [0, 3, 6], + [0, 1, 2], + ] + segment_metadata = [call[3] for call in layer.calls] + assert [metadata.causal for metadata in segment_metadata] == [True, False, True] + assert [metadata.seq_lens.tolist() for metadata in segment_metadata] == [[5, 5], [8, 8], [9, 9]] + assert [metadata.query_start_loc_cpu.tolist() for metadata in segment_metadata] == [ + [0, 2, 4], + [0, 3, 6], + [0, 1, 2], + ] + assert [metadata.positions.tolist() for metadata in segment_metadata] == [ + [3, 4, 3, 4], + [5, 6, 7, 5, 6, 7], + [8, 8], + ] + assert [metadata.slot_mappings.tolist() for metadata in segment_metadata] == [ + [[3, 4, 3, 4]], + [[5, 6, 7, 5, 6, 7]], + [[8, 8]], + ] + + +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( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)] + ) + qkv = torch.randn(1, 4, 2, 4) + metadata = SimpleNamespace( + attn_mask=torch.ones(1, 1, 4, 4, dtype=torch.bool).tril(), + full_attn_spans=[[]], + query_ranges=None, + 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"] + assert len(events) == 2 + assert events[1][1]["causal"] is True + assert events[1][1]["seq_lens"].tolist() == [4] + + +def test_omni_paged_backend_runs_unequal_hunyuan_prefixes_without_padding( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, layer, events = _make_adapter(monkeypatch, capacity=64) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=17, + seq_len=50, + kv_start_pos=33, + ), + DiffusionPagedAttentionRow( + request_id="req-1", + sequence_id=1, + query_len=17, + seq_len=54, + kv_start_pos=37, + ), + ] + ) + query = torch.arange(2 * 17 * 2 * 4, dtype=torch.float32).reshape(2, 17, 2, 4) + key = query + 100 + value = query + 200 + metadata = SimpleNamespace( + attn_mask=torch.ones(2, 1, 17, 54, dtype=torch.bool), + full_attn_spans=[[(33, 50)], [(37, 54)]], + query_ranges=None, + extra={}, + ) + + with adapter.activate(batch): + context = adapter.prepare_layer_context( + "layer-0", + query, + key, + value, + omni_attn_metadata=metadata, + ) + output = _run_omni_paged_backend(context) + + assert torch.equal(output, query) + assert batch.query_start_loc.tolist() == [0, 17, 34] + assert batch.seq_lens.tolist() == [50, 54] + assert batch.positions.tolist() == [*range(33, 50), *range(37, 54)] + assert torch.equal(context.key_write, key.reshape(34, 2, 4)) + assert torch.equal(context.value_write, value.reshape(34, 2, 4)) + assert layer.native_events == ["update", "forward"] + assert layer.updates[0][2].tolist() == batch.positions.tolist() + assert layer.calls[0][3].seq_lens.tolist() == [50, 54] + assert layer.calls[0][3].query_start_loc_cpu.tolist() == [0, 17, 34] + assert len(events) == 2 + + +def test_forward_rejects_unaligned_heterogeneous_piecewise_spans_before_cache_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4), + DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=4, seq_len=4), + ] + ) + qkv = torch.randn(2, 4, 2, 4) + metadata = SimpleNamespace( + attn_mask=torch.ones(2, 1, 4, 4, dtype=torch.bool), + full_attn_spans=[[(1, 3)], [(2, 4)]], + query_ranges=None, + extra={}, + ) + + with adapter.activate(batch), pytest.raises(ValueError, match="produce aligned segments"): + adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + + assert layer.native_events == [] + + +@pytest.mark.parametrize( + "spans", + [ + [(1, 3), (2, 4)], + [(1, 5)], + ], +) +def test_forward_rejects_invalid_piecewise_spans_before_cache_write( + monkeypatch: pytest.MonkeyPatch, + spans: list[tuple[int, int]], +) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + 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=[spans], extra={}) + + with adapter.activate(batch), pytest.raises(ValueError, match="sorted, non-overlapping, and within the sequence"): + adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + + assert layer.native_events == [] + + +def test_forward_rejects_non_4d_mask_with_piecewise_spans_before_cache_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + 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( + attn_mask=torch.ones(1, 4, 4, dtype=torch.bool), + full_attn_spans=[[(1, 3)]], + extra={}, + ) + + with adapter.activate(batch), pytest.raises(NotImplementedError, match="requires a 4D tensor"): + adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + + assert layer.native_events == [] + + +def test_forward_rejects_piecewise_query_ranges_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + 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( + attn_mask=torch.ones(1, 1, 4, 4, dtype=torch.bool), + full_attn_spans=[[(1, 3)]], + query_ranges=(object(),), + extra={}, + ) + + with adapter.activate(batch), pytest.raises(NotImplementedError, match="query_ranges"): + adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + + assert layer.native_events == [] + + +def test_forward_rejects_piecewise_non_suffix_query_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=2, + seq_len=6, + kv_start_pos=2, + ) + ] + ) + qkv = torch.randn(1, 2, 2, 4) + metadata = SimpleNamespace(full_attn_spans=[[(2, 4)]], extra={}) + + with adapter.activate(batch), pytest.raises(ValueError, match="requires each query/write span to end"): + adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + + assert layer.native_events == [] + + +def test_forward_accepts_non_aligned_current_write(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=2, + seq_len=5, + kv_start_pos=3, + ) + ] + ) + query = torch.randn(1, 2, 2, 4) + key = torch.arange(1 * 2 * 2 * 4, dtype=torch.float32).reshape(1, 2, 2, 4) + value = key + 100 + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", query, key, value) + + assert context.query.shape == (2, 2, 4) + assert torch.equal(context.key_write, key.reshape(2, 2, 4)) + assert torch.equal(context.value_write, value.reshape(2, 2, 4)) + assert context.slot_mapping.tolist() == [3, 4] + assert layer.native_events == [] + + +def test_forward_rejects_untranslated_omni_metadata_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)] + ) + qkv = torch.randn(1, 2, 4) + metadata = SimpleNamespace(attn_mask=torch.ones(1, 1), extra={}) + + with adapter.activate(batch), pytest.raises(NotImplementedError, match="cannot translate"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv, omni_attn_metadata=metadata) + + assert layer.updates == [] + + +def test_forward_accepts_flattened_hidden_dimension(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)] + ) + query = torch.randn(3, 8) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", query, query, query) + + assert context.query.shape == (3, 2, 4) + assert context.restore_output(context.query).shape == query.shape + + +def test_forward_accepts_one_flattened_token_for_single_head(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch, num_heads=1, num_kv_heads=1) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)] + ) + qkv = torch.randn(1, 4) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + output = _run_omni_paged_backend(context) + + assert context.query.shape == (1, 1, 4) + assert context.query_token_shape == (1,) + assert not context.query_has_head_dims + assert output.shape == qkv.shape + + +def test_forward_requires_active_prepared_batch(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + qkv = torch.randn(1, 2, 4) + + with pytest.raises(RuntimeError, match="adapter.activate"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + + +def test_new_preparation_invalidates_native_buffer_views_in_older_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + old_batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)] + ) + adapter.prepare_batch([DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1)]) + + with pytest.raises(ValueError, match="stale"): + with adapter.activate(old_batch): + pass + + +def test_block_table_change_invalidates_prepared_native_buffer_views( + 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)] + ) + + adapter.invalidate_prepared_batches() + + with pytest.raises(ValueError, match="stale"): + with adapter.activate(batch): + pass + + +def test_block_table_change_is_rejected_during_active_forward(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)] + ) + + with adapter.activate(batch), pytest.raises(RuntimeError, match="during an active forward"): + adapter.invalidate_prepared_batches() + + +def test_prepare_batch_rejects_active_native_buffers(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)] + ) + + with adapter.activate(batch), pytest.raises(RuntimeError, match="during an active forward"): + adapter.prepare_batch([DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1)]) + + +def test_forward_rejects_tensor_count_different_from_prepared_batch(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)] + ) + qkv = torch.randn(2, 2, 4) + + with adapter.activate(batch), pytest.raises(ValueError, match="token count"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + + +def test_forward_rejects_dtype_different_from_model_activation_dtype(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)] + ) + qkv = torch.randn(1, 2, 4, dtype=torch.float64) + + with adapter.activate(batch), pytest.raises(ValueError, match="model activation dtype"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + + +def test_forward_rejects_invalid_key_token_count(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)] + ) + query = torch.randn(4, 2, 4) + key = torch.randn(3, 2, 4) + + with adapter.activate(batch), pytest.raises(ValueError, match="key token count"): + adapter.prepare_layer_context("layer-0", query, key, query) + + +def test_forward_rejects_batched_layout_that_does_not_match_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.randn(2, 2, 2, 4) + + with adapter.activate(batch), pytest.raises(ValueError, match="must match prepared rows"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + + +def test_forward_rejects_padded_full_kv(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=3, + kv_start_pos=2, + ), + DiffusionPagedAttentionRow( + request_id="req-1", + sequence_id=1, + query_len=1, + seq_len=5, + kv_start_pos=4, + ), + ] + ) + query = torch.randn(2, 1, 2, 4) + full_kv = torch.randn(2, 4, 2, 4) + + with adapter.activate(batch), pytest.raises(ValueError, match="batched key layout must match prepared rows"): + adapter.prepare_layer_context("layer-0", query, full_kv, full_kv) + + +def test_causal_batch_rejects_non_suffix_write(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch, non_causal=False) + row = DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=2, + seq_len=8, + kv_start_pos=3, + ) + + batch = adapter.prepare_batch([row]) + qkv = torch.randn(2, 2, 4) + + with adapter.activate(batch), pytest.raises(ValueError, match="requires each query/write span to end"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + + +def test_prepare_batch_rejects_native_decode_after_prefill_order(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, block_tables, _, _ = _make_adapter( + monkeypatch, + non_causal=False, + reorder_batch_threshold=1, + ) + rows = [ + DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=2, seq_len=2), + DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1), + ] + + with pytest.raises(ValueError, match="decode/short-query rows before"): + adapter.prepare_batch(rows) + + assert block_tables.gather_calls == [] + assert block_tables.slot_calls == [] + + +def test_prepare_batch_rejects_seq_len_beyond_installed_row_capacity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, block_tables, _, _ = _make_adapter(monkeypatch) + block_tables.num_blocks.np[0, 2] = 1 + row = DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=5, + seq_len=5, + ) + + with pytest.raises(ValueError, match="installed Worker row contains only 1 blocks"): + adapter.prepare_batch([row]) + + assert block_tables.gather_calls == [] + assert block_tables.slot_calls == [] + + +def test_prepare_batch_rejects_seq_len_beyond_installed_logical_length( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, block_tables, _, _ = _make_adapter(monkeypatch) + adapter.resolve_row = lambda *_args: DiffusionPagedAttentionRowBinding(2, 4) + row = DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=5, + seq_len=5, + ) + + with pytest.raises(ValueError, match="installed allocation has logical length 4"): + adapter.prepare_batch([row]) + + assert block_tables.gather_calls == [] + assert block_tables.slot_calls == [] + + +def test_prepare_batch_passes_dcp_local_seq_lens_to_native_builder( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, block_tables, _, events = _make_adapter(monkeypatch) + block_tables.cp_size = 2 + block_tables.cp_rank = 1 + block_tables.cp_interleave = 2 + block_tables.num_blocks.np[0, 2] = 1 + + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)] + ) + + assert batch.row_indices.tolist() == [2] + assert events[0][1]["dcp_local_seq_lens"].tolist() == [2] + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"sequence_id": 0, "context_id": "text"}, "exactly one"), + ({}, "exactly one"), + ({"sequence_id": 0, "query_len": 0}, "query_len"), + ({"sequence_id": 0, "kv_start_pos": 4, "query_len": 2, "seq_len": 5}, "exceeds seq_len"), + ], +) +def test_row_contract_rejects_invalid_identity_or_span(kwargs: dict, message: str) -> None: + values = { + "request_id": "req-0", + "query_len": 1, + "seq_len": 1, + **kwargs, + } + with pytest.raises(ValueError, match=message): + DiffusionPagedAttentionRow(**values) + + +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())) + native_backend = SimpleNamespace( + get_name=lambda: "ASCEND", + indexes_kv_by_block_stride=lambda: True, + get_impl_cls=lambda: impl_cls, + forward_includes_kv_cache_update=True, + ) + + original_backend_per_kind = {"full": object()} + config = SimpleNamespace( + attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind), + 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)) + return native_backend + + monkeypatch.setattr(adapter_module, "get_attn_backend", select_backend) + monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext()) + layer = SimpleNamespace(num_heads=8, softmax_scale=0.125) + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=8, + dtype=torch.float16, + non_causal=True, + ) + + native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter( + layer_name="layer-0", + layer=layer, + spec=spec, + vllm_config=config, + device=torch.device("cpu"), + ulysses_degree=2, + ) + + assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {})] + assert config.attention_config.backend is None + assert config.attention_config.backend_per_kind is original_backend_per_kind + assert native_layer.num_heads == 4 + assert native_layer.num_kv_heads == 2 + assert native_layer.spec.num_kv_heads == 2 + assert native_layer.spec.indexes_kv_by_block_stride is True + assert native_layer._q_scale_float == 1.0 + assert native_layer._k_scale_float == 1.0 + assert native_layer._v_scale_float == 1.0 + impl_kwargs = impl_cls.call_args.kwargs + assert impl_kwargs["num_heads"] == 4 + assert impl_kwargs["num_kv_heads"] == 2 + + +def test_omni_attention_wraps_paged_kernel_with_sp_hooks() -> None: + events: list[str] = [] + + class Strategy: + name = "ulysses" + + def pre_attention(self, query, key, value, metadata): + events.append("pre") + return query + 1, key + 2, value + 3, metadata, object() + + def post_attention(self, output, _ctx): + events.append("post") + return output + 4 + + class Adapter: + def prepare_layer_context(self, layer_name, query, key, value, *, omni_attn_metadata): + events.append("prepare") + assert layer_name == "layer-0" + assert omni_attn_metadata is None + assert torch.equal(key, original + 2) + assert torch.equal(value, original + 3) + return SimpleNamespace(query=query) + + class Backend: + def forward_paged(self, context): + events.append("kernel") + return context.query * 2 + + 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() + layer._with_kv_cache_dtype = lambda metadata: metadata + original = torch.zeros(1, 2, 2, 4) + + with set_forward_context(), override_paged_kv_adapter(Adapter()): + assert layer.is_paged_kv_active() + output = layer._forward_impl(original, original, original) + + assert events == ["pre", "prepare", "kernel", "post"] + assert torch.equal(output, torch.full_like(original, 6)) + + +def test_omni_attention_rejects_paged_request_for_backend_without_paged_support() -> None: + class Adapter: + def prepare_layer_context(self, *args, **kwargs): + raise AssertionError("an unsupported backend must fail before adapter preparation") + + 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=False, get_name=lambda: "SDPA") + layer.use_ring = False + layer._no_parallel_strategy = object() + layer._get_active_parallel_strategy = lambda: layer._no_parallel_strategy + + with ( + set_forward_context(), + override_paged_kv_adapter(Adapter()), + pytest.raises(NotImplementedError, match="backend with paged support"), + ): + layer._forward_impl(torch.zeros(1, 2, 2, 4), torch.zeros(1, 2, 2, 4), torch.zeros(1, 2, 2, 4)) + + +def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None: + events: list[str] = [] + + class Strategy: + name = "ulysses" + + def pre_attention(self, query, key, value, metadata): + events.append("pre") + return query, key, value, metadata, object() + + def post_attention(self, output, _ctx): + events.append("post") + return output + + layer = Attention.__new__(Attention) + nn.Module.__init__(layer) + layer.paged_kv_cache_role = "primary" + layer.use_ring = False + layer._no_parallel_strategy = object() + layer._get_active_parallel_strategy = lambda: Strategy() + layer._with_kv_cache_dtype = lambda metadata: metadata + layer._run_local_attention = lambda query, _key, _value, _metadata: events.append("dense") or query + qkv = torch.zeros(1, 2, 2, 4) + + assert not layer.is_paged_kv_active() + output = layer._forward_impl(qkv, qkv, qkv) + + assert events == ["pre", "dense", "post"] + assert output is 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 new file mode 100644 index 000000000..51ab6567d --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from torch.nn import functional as F +from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache +from vllm.v1.worker.gpu.block_table import BlockTables + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl +from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( + DiffusionPagedAttentionAdapter, + DiffusionPagedAttentionLayerAdapter, + DiffusionPagedAttentionRow, + DiffusionPagedAttentionRowBinding, +) +from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] + +_LAYER_NAME = "model.layers.0.attn" +_NUM_HEADS = 2 +_HEAD_SIZE = 64 +_BLOCK_SIZE = 16 + + +class _SmokeDiffusionAttention(nn.Module): + def __init__(self) -> None: + super().__init__() + self.num_heads = _NUM_HEADS + self.num_kv_heads = _NUM_HEADS + self.head_size = _HEAD_SIZE + self.softmax_scale = _HEAD_SIZE**-0.5 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention") +def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> None: + device = torch.device("cuda", torch.accelerator.current_device_index()) + vllm_config = VllmConfig( + cache_config=CacheConfig( + block_size=_BLOCK_SIZE, + cache_dtype="float16", + enable_prefix_caching=False, + ), + attention_config=AttentionConfig(use_non_causal=True), + ) + model_config = _DiffusionVllmModelConfig( + model=None, + dtype=torch.float16, + max_model_len=32, + original_max_model_len=32, + hf_config=SimpleNamespace(model_type="diffusion_paged_attention_smoke"), + ) + vllm_config.model_config = model_config + model_config.set_attention_geometry( + num_heads=_NUM_HEADS, + num_kv_heads=_NUM_HEADS, + head_size=_HEAD_SIZE, + ) + vllm_config.compilation_config.static_forward_context.clear() + + diffusion_layer = _SmokeDiffusionAttention().to(device) + spec = FullAttentionSpec( + block_size=_BLOCK_SIZE, + num_kv_heads=_NUM_HEADS, + head_size=_HEAD_SIZE, + dtype=torch.float16, + non_causal=True, + ) + with set_current_vllm_config(vllm_config): + native_layer = DiffusionPagedAttentionLayerAdapter( + layer_name=_LAYER_NAME, + layer=diffusion_layer, + spec=spec, + vllm_config=vllm_config, + device=device, + ) + vllm_config.compilation_config.static_forward_context[_LAYER_NAME] = native_layer + + canonical_spec = native_layer.spec + kv_cache_config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=canonical_spec.page_size_bytes * 4, + shared_by=[_LAYER_NAME], + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec( + layer_names=[_LAYER_NAME], + kv_cache_spec=canonical_spec, + ) + ], + ) + with set_current_vllm_config(vllm_config): + attn_groups, _, kernel_block_sizes = init_attn_backend( + kv_cache_config, + vllm_config, + device, + ) + runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = [] + init_kv_cache( + runner_kv_caches, + vllm_config.compilation_config.static_forward_context, + kv_cache_config, + attn_groups, + device, + vllm_config.cache_config.cache_dtype, + kernel_block_sizes, + vllm_config, + ) + block_tables = BlockTables( + block_sizes=[_BLOCK_SIZE], + max_num_reqs=1, + max_num_batched_tokens=32, + max_num_blocks_per_group=[2], + device=device, + kernel_block_sizes=kernel_block_sizes, + ) + + block_tables.append_block_ids(0, ([3, 1],), overwrite=True) + block_tables.apply_staged_writes() + adapter = DiffusionPagedAttentionAdapter( + vllm_config=vllm_config, + device=device, + kv_cache_config=kv_cache_config, + block_tables=block_tables, + attn_groups=attn_groups, + layers={_LAYER_NAME: native_layer}, + resolve_row=lambda _request_id, _sequence_id, _context_id: DiffusionPagedAttentionRowBinding( + row_index=0, + max_seq_len=19, + ), + ) + prefix_batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=17, + seq_len=17, + ) + ] + ) + + 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) + omni_backend = FlashAttentionImpl( + num_heads=_NUM_HEADS, + head_size=_HEAD_SIZE, + softmax_scale=diffusion_layer.softmax_scale, + num_kv_heads=_NUM_HEADS, + ) + with adapter.activate(prefix_batch): + prefix_context = adapter.prepare_layer_context(_LAYER_NAME, query[:17], key[:17], value[:17]) + prefix_output = omni_backend.forward_paged(prefix_context) + prefix_slot_mappings = prefix_batch.slot_mappings.clone() + + suffix_batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=2, + seq_len=19, + kv_start_pos=17, + ) + ] + ) + with adapter.activate(suffix_batch): + suffix_context = adapter.prepare_layer_context(_LAYER_NAME, query[17:], key[17:], value[17:]) + suffix_output = omni_backend.forward_paged(suffix_context) + suffix_slot_mappings = suffix_batch.slot_mappings.clone() + + piecewise_batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=19, + seq_len=19, + ) + ] + ) + mixed_mask = torch.ones(19, 19, dtype=torch.bool, device=device).tril() + mixed_mask[5:10, 5:10] = True + mixed_metadata = AttentionMetadata( + attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0), + full_attn_spans=[[(5, 10)]], + ) + with adapter.activate(piecewise_batch): + piecewise_context = adapter.prepare_layer_context( + _LAYER_NAME, + query, + key, + value, + omni_attn_metadata=mixed_metadata, + ) + piecewise_output = omni_backend.forward_paged(piecewise_context) + + prefix_reference = ( + F.scaled_dot_product_attention( + query[:17].transpose(0, 1).unsqueeze(0), + key[:17].transpose(0, 1).unsqueeze(0), + value[:17].transpose(0, 1).unsqueeze(0), + is_causal=False, + scale=diffusion_layer.softmax_scale, + ) + .squeeze(0) + .transpose(0, 1) + ) + suffix_reference = ( + F.scaled_dot_product_attention( + query[17:].transpose(0, 1).unsqueeze(0), + key.transpose(0, 1).unsqueeze(0), + value.transpose(0, 1).unsqueeze(0), + is_causal=False, + scale=diffusion_layer.softmax_scale, + ) + .squeeze(0) + .transpose(0, 1) + ) + piecewise_reference = ( + F.scaled_dot_product_attention( + query.transpose(0, 1).unsqueeze(0), + key.transpose(0, 1).unsqueeze(0), + value.transpose(0, 1).unsqueeze(0), + attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0), + scale=diffusion_layer.softmax_scale, + ) + .squeeze(0) + .transpose(0, 1) + ) + torch.accelerator.synchronize(device) + + assert prefix_slot_mappings.cpu().tolist() == [list(range(3 * _BLOCK_SIZE, 4 * _BLOCK_SIZE)) + [_BLOCK_SIZE]] + assert suffix_slot_mappings.cpu().tolist() == [[_BLOCK_SIZE + 1, _BLOCK_SIZE + 2]] + torch.testing.assert_close(prefix_output, prefix_reference, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(suffix_output, suffix_reference, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(piecewise_output, piecewise_reference, rtol=2e-2, atol=2e-2) diff --git a/tests/diffusion/diffusion_kv/test_spec_discovery.py b/tests/diffusion/diffusion_kv/test_spec_discovery.py index 1d1d6c8d9..7c6d94a0e 100644 --- a/tests/diffusion/diffusion_kv/test_spec_discovery.py +++ b/tests/diffusion/diffusion_kv/test_spec_discovery.py @@ -25,6 +25,23 @@ class _Backend: return True +class _RunnerKVBackend: + def __init__(self) -> None: + self.kv_cache_config = None + + def register_kv_cache_layers(self, layers): + return {layer_name: spec for layer_name, (_layer, spec) in layers.items()} + + def initialize_kv_cache(self, config) -> None: + configured_layers = {layer_name for group in config.kv_cache_groups for layer_name in group.layer_names} + if configured_layers != {"image_attention"}: + raise ValueError( + "Rank-local Diffusion KVCacheConfig layer mismatch: " + f"expected=['image_attention'], configured={sorted(configured_layers)}" + ) + self.kv_cache_config = config + + def _attention(*, enabled: bool) -> Attention: attention = Attention.__new__(Attention) nn.Module.__init__(attention) @@ -47,6 +64,7 @@ def _runner(attention: Attention) -> DiffusionModelRunner: pipeline = nn.Module() pipeline.image_attention = attention runner.pipeline = pipeline + runner.diffusion_kv_backend = _RunnerKVBackend() return runner @@ -102,14 +120,21 @@ def test_worker_selects_its_rank_local_config() -> None: worker = object.__new__(DiffusionWorker) worker.rank = 1 worker.od_config = SimpleNamespace(num_gpus=2) + worker.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=None), + cache_config=SimpleNamespace(num_gpu_blocks=None), + ) + worker._maybe_get_memory_pool_context = lambda _tag: nullcontext() worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config)) - configs = [object(), object()] + configs = [SimpleNamespace(num_blocks=4), SimpleNamespace(num_blocks=8)] - worker.set_kv_cache_configs(configs) + worker.set_kv_cache_configs(configs, 64) assert worker.installed is configs[1] + assert worker.vllm_config.model_config.max_model_len == 64 + assert worker.vllm_config.cache_config.num_gpu_blocks == 8 with pytest.raises(ValueError, match="rank count mismatch"): - worker.set_kv_cache_configs(configs[:1]) + worker.set_kv_cache_configs(configs[:1], 64) def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None: 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 fed56af2d..98026564c 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 @@ -19,8 +19,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] _TRANSFORMER_MODULE = "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer" -NUM_HEADS = 4 -NUM_KV_HEADS = 2 +NUM_HEADS = 32 +NUM_KV_HEADS = 8 HEAD_DIM = 16 IMAGE_TOKEN_LEN = 8 SUFFIX_TOKEN_LEN = 3 @@ -35,8 +35,14 @@ SCALING = 1.0 / math.sqrt(HEAD_DIM) class MockAttention(nn.Module): def __init__(self, num_heads, head_size, causal=False, softmax_scale=None, num_kv_heads=None, **kwargs): super().__init__() + self.paged_kv_active = False + self.calls = [] + + 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)) return query @@ -126,6 +132,37 @@ def test_cache_manager_registers_attention_without_adding_dense_state() -> None: assert mgr.state_dict() == {} +@pytest.mark.parametrize( + ("paged_kv_active", "expected_kv_heads"), + [(False, NUM_HEADS), (True, NUM_KV_HEADS)], +) +def test_gqa_kv_stays_compressed_for_paged_attention( + paged_kv_active: bool, + expected_kv_heads: int, +) -> None: + mgr = _make_cache_mgr() + mgr.attn.paged_kv_active = paged_kv_active + q_len = 3 + key, value = _make_known_kv(q_len) + + _call_mgr( + mgr, + bs=1, + q_len=q_len, + seq_len=q_len, + key_flat=key, + value_flat=value, + first_step=True, + num_image_tokens=0, + gen_timestep_scatter_index=_gen_timestep_index(1, 0), + ) + + _, key_input, value_input = mgr.attn.calls[-1] + expected_shape = (1, q_len, expected_kv_heads, HEAD_DIM) + assert key_input.shape == expected_shape + assert value_input.shape == expected_shape + + # ============================================================ # Test 1: No AR KV — basic cache → reuse # ============================================================ diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index dfe4a41e8..60d0c63b8 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -149,10 +149,18 @@ class TestCreateDefaultDiffusion: assert ea["enforce_eager"] is True assert ea["lora_path"] == "/tmp/lora" - def test_diffusion_kv_mode_roundtrip(self): - od = _roundtrip_diffusion_config(model="x", diffusion_kv_mode="paged_scheduler") + def test_diffusion_kv_mode_roundtrip(self, monkeypatch): + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) + od = _roundtrip_diffusion_config( + model="x", + diffusion_kv_mode="paged_scheduler", + diffusion_kv_max_rows_per_request=2, + ) assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + assert od.diffusion_kv_max_rows_per_request == 2 def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch): monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500) diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index b626f8ac6..ba7d7d5d2 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -12,6 +12,7 @@ import pytest from vllm_omni.diffusion import diffusion_engine as diffusion_engine_module from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -88,6 +89,50 @@ def test_emit_step_outputs_finalizes_finished_request_without_stream() -> None: assert engine.scheduler.get_request_state(request_id) is None +def _make_kv_cleanup_engine(mode: DiffusionKVCacheMode) -> tuple[DiffusionEngine, Mock]: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace(diffusion_kv_mode=mode) + cleanup = Mock() + engine.executor = SimpleNamespace(remove_diffusion_kv_requests=cleanup) + return engine, cleanup + + +def test_paged_terminal_output_clears_worker_rows() -> None: + engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER) + engine._finalize_finished_request = lambda request_id, *_args: request_id + engine._put_output = lambda *_args: None + + engine._emit_finished_outputs({"req-0", "req-1"}) + + cleanup.assert_called_once() + assert set(cleanup.call_args.args[0]) == {"req-0", "req-1"} + + +def test_paged_abort_clears_worker_rows_before_scheduler_free() -> None: + engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER) + engine.scheduler = SimpleNamespace( + get_request_state=lambda _request_id: None, + finish_requests=Mock(side_effect=AssertionError("unknown request must not be logically freed")), + ) + + engine._abort_requests(["req-idle", "req-idle"]) + + cleanup.assert_called_once_with(["req-idle"]) + engine.scheduler.finish_requests.assert_not_called() + + +def test_dense_terminal_and_abort_paths_skip_worker_row_cleanup() -> None: + engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.DENSE_LEGACY) + engine._finalize_finished_request = lambda request_id, *_args: request_id + engine._put_output = lambda *_args: None + engine.scheduler = SimpleNamespace(get_request_state=lambda _request_id: None) + + engine._emit_finished_outputs({"req-0"}) + engine._abort_requests(["req-idle"]) + + cleanup.assert_not_called() + + def test_init_accepts_custom_scheduler(monkeypatch: pytest.MonkeyPatch) -> None: od_config = SimpleNamespace( custom_pipeline_args=None, diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 66da14fd5..c331f08ee 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -11,6 +11,7 @@ import torch import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module from tests.helpers.mark import hardware_test from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request @@ -193,6 +194,11 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar parallel_config=SimpleNamespace(use_hsdp=False), streaming_output=False, ) + runner.diffusion_kv_backend = DiffusionKVModelRunnerBackend( + vllm_config=runner.vllm_config, + od_config=runner.od_config, + device=runner.device, + ) runner.kv_transfer_manager = SimpleNamespace( receive_kv_cache=lambda req, target_device=None: None, receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None, diff --git a/tests/diffusion/test_diffusion_plugin_hooks.py b/tests/diffusion/test_diffusion_plugin_hooks.py index a437a572a..1885d3338 100644 --- a/tests/diffusion/test_diffusion_plugin_hooks.py +++ b/tests/diffusion/test_diffusion_plugin_hooks.py @@ -40,6 +40,12 @@ class TestPlatformDiffusionHooks: result = OmniPlatform.get_diffusion_model_runner_cls() assert result == "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner" + def test_get_diffusion_kv_block_tables_cls_default(self): + """Test the default paged-KV BlockTables implementation.""" + from vllm.v1.worker.gpu.block_table import BlockTables + + assert OmniPlatform.get_diffusion_kv_block_tables_cls() is BlockTables + def test_oot_enum_exists(self): """Test OOT is a valid platform enum value.""" assert OmniPlatformEnum.OOT.value == "oot" diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 0964833d7..96f8249c5 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -17,6 +17,8 @@ import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module from tests.helpers.mark import hardware_test from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D from vllm_omni.diffusion.distributed.parallel_state import ( @@ -306,11 +308,14 @@ def _make_vllm_config(): ) -def _make_runner(): +def _make_runner( + diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY, +): runner = object.__new__(DiffusionModelRunner) runner.vllm_config = _make_vllm_config() runner.od_config = SimpleNamespace( cache_backend=None, + diffusion_kv_mode=diffusion_kv_mode, parallel_config=SimpleNamespace(use_hsdp=False), streaming_output=False, ) @@ -330,6 +335,7 @@ def _make_distributed_runner(mode: str, device: torch.device): runner.vllm_config = _make_vllm_config() runner.od_config = SimpleNamespace( cache_backend=None, + diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY, parallel_config=SimpleNamespace(use_hsdp=False), streaming_output=False, ) @@ -502,6 +508,81 @@ def test_step_profiler_reports_denoise_step_as_diffuse(): class TestRunner: """DiffusionModelRunner.execute_stepwise""" + @pytest.mark.parametrize( + ("diffusion_kv_mode", "should_cleanup"), + [ + (DiffusionKVCacheMode.DENSE_LEGACY, False), + (DiffusionKVCacheMode.PAGED_SCHEDULER, True), + ], + ) + def test_scheduler_finished_cleanup_only_runs_for_paged_kv( + self, + mocker: MockerFixture, + diffusion_kv_mode: DiffusionKVCacheMode, + should_cleanup: bool, + ): + runner = _make_runner(diffusion_kv_mode) + cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests") + expected_output = object() + mocker.patch.object(runner, "_execute_stepwise_core", return_value=expected_output) + scheduler_output = DiffusionSchedulerOutput( + step_id=1, + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + finished_req_ids={"finished-request"}, + num_running_reqs=0, + num_waiting_reqs=0, + ) + + output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + + assert output is expected_output + if should_cleanup: + cleanup.assert_called_once_with(["finished-request"]) + else: + cleanup.assert_not_called() + + @pytest.mark.parametrize( + ("diffusion_kv_mode", "should_cleanup"), + [ + (DiffusionKVCacheMode.DENSE_LEGACY, False), + (DiffusionKVCacheMode.PAGED_SCHEDULER, True), + ], + ) + def test_terminal_cleanup_only_runs_for_paged_kv( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + diffusion_kv_mode: DiffusionKVCacheMode, + should_cleanup: bool, + ): + runner = _make_runner(diffusion_kv_mode) + cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests") + install = mocker.patch.object(runner, "install_diffusion_kv_metadata") + req = _make_step_request() + scheduler_output = _make_scheduler_output(req, step_id=0) + if should_cleanup: + scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata = DiffusionKVMetadata( + request_id=req.request_id, + allocation_generation=1, + sequences=(), + ) + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + + first_output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + assert first_output.get_request_output(req.request_id).finished is False + cleanup.assert_not_called() + + output = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1)) + + assert output.get_request_output(req.request_id).finished is True + if should_cleanup: + install.assert_called_once_with(scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata) + cleanup.assert_called_once_with([req.request_id]) + else: + install.assert_not_called() + cleanup.assert_not_called() + def test_completes_request_and_clears_state(self, monkeypatch): runner = _make_runner() req = _make_step_request() diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index c1e24ea4b..b60f574a1 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -678,6 +678,7 @@ class _DiffusionConfigProjection: cache_config: Any = field(default_factory=dict) enable_cache_dit_summary: bool = False diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY + diffusion_kv_max_rows_per_request: int | None = Field(default=None, ge=1, strict=True) enable_prompt_embed_cache: bool = False prompt_embed_cache_size: int = Field(default=32, ge=1) enable_session_state_manager: bool = False @@ -816,6 +817,11 @@ class _DiffusionConfigProjection: ) self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode) + if ( + self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + and self.diffusion_kv_max_rows_per_request is None + ): + raise ValueError("paged_scheduler requires diffusion_kv_max_rows_per_request to be set") self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps) self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers) diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 4a34aa765..95f0f2736 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -15,6 +15,11 @@ class AttentionBackend(ABC): accept_output_buffer: bool = False supports_piecewise_spans: bool = False + # A backend that supports this capability can consume the opaque paged-KV + # context prepared by the diffusion Worker data plane. Keeping the + # capability on the backend class prevents a paged request from silently + # falling back to dense attention on an incompatible implementation. + supports_paged_kv: bool = False # The backend can represent a contiguous valid K/V prefix by slicing the # tensors instead of materializing a padding mask. Models may use this to # avoid a slower masked-attention plan when tail padding is not semantic. @@ -221,6 +226,18 @@ class AttentionImpl(ABC, Generic[T]): else: raise NotImplementedError(f"No forward implementation for platform: {current_omni_platform}") + def forward_paged(self, paged_kv_context: Any) -> torch.Tensor: + """Execute one Worker-prepared paged-KV attention call. + + The context is intentionally opaque to the common attention layer. + Backends opt in by setting ``supports_paged_kv`` on their backend + class and implementing this method. Dense callers continue to use + ``forward`` unchanged. + """ + + del paged_kv_context + raise NotImplementedError(f"{type(self).__name__} does not support paged KV attention") + def forward_cuda( self, query: torch.Tensor, diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index bdef1dd65..d287ef7d2 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -15,6 +15,7 @@ from vllm_omni.diffusion.attention.backends.abstract import ( from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( piecewise_attn, + run_paged_piecewise_plan, ) from vllm_omni.diffusion.config import get_current_diffusion_config_or_none from vllm_omni.platforms import current_omni_platform @@ -25,6 +26,7 @@ logger = init_logger(__name__) class FlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True supports_piecewise_spans: bool = True + supports_paged_kv: bool = True @classmethod def supports_packed_mask_free(cls) -> bool: @@ -246,6 +248,73 @@ class FlashAttentionImpl(AttentionImpl): # (b s) h d -> b s h d return out.reshape(batch_size, q_len, *out.shape[1:]) + def forward_paged(self, paged_kv_context) -> torch.Tensor: + """Run platform-native paged attention through the Omni backend. + + ``DiffusionPagedAttentionAdapter`` owns BlockTables and prepares the + 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. + """ + + layer = getattr(paged_kv_context, "layer", None) + if layer is None: + raise TypeError("paged_kv_context must expose a native diffusion attention layer") + kv_cache = getattr(layer, "kv_cache", None) + if kv_cache is None: + raise RuntimeError(f"Native KV cache is not bound for diffusion layer {layer.layer_name!r}") + 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( + layer, + paged_kv_context.key_write, + paged_kv_context.value_write, + kv_cache, + paged_kv_context.slot_mapping, + ) + + def run_native_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + native_metadata, + ) -> torch.Tensor: + 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, + kv_cache, + native_metadata, + output, + ) + + if paged_kv_context.piecewise_plan is not None: + output = run_paged_piecewise_plan( + paged_kv_context.query, + paged_kv_context.key_write, + paged_kv_context.value_write, + paged_kv_context.piecewise_plan, + paged_kv_context.piecewise_native_metadata, + run_native_attention, + ) + else: + output = run_native_attention( + paged_kv_context.query, + paged_kv_context.key_write, + paged_kv_context.value_write, + paged_kv_context.native_metadata, + ) + return paged_kv_context.restore_output(output) + def forward_cuda( self, query: torch.Tensor, diff --git a/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py b/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py index 5d83d8dbf..caa942987 100644 --- a/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py +++ b/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py @@ -15,6 +15,8 @@ Per segment: from __future__ import annotations +from collections.abc import Callable, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Literal, NamedTuple import torch @@ -123,3 +125,109 @@ def piecewise_attn( if len(outputs) == 1: return outputs[0] return torch.cat(outputs, dim=1) + + +@dataclass(frozen=True, slots=True) +class PagedPiecewiseSegment: + """One aligned segment across rows in a packed paged batch.""" + + row_segments: tuple[Segment, ...] + query_indices: torch.Tensor + + +@dataclass(frozen=True, slots=True) +class PagedPiecewisePlan: + """Piecewise segment batches for one packed paged-attention batch.""" + + spans: tuple[tuple[tuple[int, int], ...], ...] + segments: tuple[PagedPiecewiseSegment, ...] + num_query_tokens: int + + +PagedPiecewiseRunner = Callable[ + [torch.Tensor, torch.Tensor, torch.Tensor, object], + torch.Tensor, +] + + +def build_paged_piecewise_plan( + full_attn_spans: Sequence[Sequence[tuple[int, int]]], + query_offsets: Sequence[int], + query_lens: Sequence[int], + seq_lens: Sequence[int], + *, + device: torch.device | str | None = None, +) -> PagedPiecewisePlan: + """Build packed indices for corresponding piecewise segments in each row.""" + + row_count = len(full_attn_spans) + 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) + + packed_offsets = [0] + for row_spans, query_offset, query_len, seq_len in zip(spans, query_offsets, query_lens, seq_lens, strict=True): + previous_end = 0 + for start, end in row_spans: + if start < previous_end or not 0 <= start < end <= seq_len: + raise ValueError("Paged piecewise spans must be sorted, non-overlapping, and within the sequence") + previous_end = end + if query_offset < 0 or query_len <= 0 or query_offset + query_len > seq_len: + raise ValueError("Paged piecewise query range must be within the sequence") + packed_offsets.append(packed_offsets[-1] + query_len) + + segments_by_row = tuple( + tuple(build_segments(row_spans, query_offset, query_len)) + for row_spans, query_offset, query_len in zip(spans, query_offsets, query_lens, strict=True) + ) + if len({len(row_segments) for row_segments in segments_by_row}) != 1: + raise ValueError("Paged piecewise rows must produce aligned segments") + + packed_segments = [] + 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 = [] + 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 + query_indices.extend(range(packed_offsets[row_index] + local_start, packed_offsets[row_index] + local_end)) + packed_segments.append( + PagedPiecewiseSegment( + row_segments=tuple(row_segments), + query_indices=torch.tensor(query_indices, dtype=torch.long, device=device), + ) + ) + + return PagedPiecewisePlan(spans, tuple(packed_segments), packed_offsets[-1]) + + +def run_paged_piecewise_plan( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + plan: PagedPiecewisePlan, + segment_metadata: Sequence[object], + segment_runner: PagedPiecewiseRunner, +) -> torch.Tensor: + """Run and scatter piecewise attention over packed Q/K/V write tensors.""" + + 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]}") + + output = None + for segment, metadata in zip(plan.segments, segment_metadata, strict=True): + indices = segment.query_indices + segment_output = segment_runner( + query.index_select(0, indices), + key.index_select(0, indices), + value.index_select(0, indices), + metadata, + ) + if output is None: + output = segment_output.new_empty((query.shape[0], *segment_output.shape[1:])) + output.index_copy_(0, indices, segment_output) + + assert output is not None + return output diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index ad9fe2b36..2ab4c6f91 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -24,7 +24,11 @@ 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.distributed.parallel_state import get_sp_group -from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available +from vllm_omni.diffusion.forward_context import ( + get_forward_context, + get_ulysses_mode, + is_forward_context_available, +) from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -73,6 +77,11 @@ class Attention(nn.Module): self.role = role self.role_category = role_category self.qkv_layout = qkv_layout + # ``prefix`` is also the stable layer identity used by vLLM's native + # KV-cache metadata. Keep it on the Omni layer so the active paged + # adapter can dispatch the already-resharded Q/K/V to the matching + # native cache tensor without replacing this Omni execution path. + self.prefix = prefix if paged_kv_cache_role == "": raise ValueError("paged_kv_cache_role must be non-empty when provided") self.paged_kv_cache_role = paged_kv_cache_role @@ -311,19 +320,48 @@ class Attention(nn.Module): ) -> torch.Tensor: # Get the appropriate parallel strategy based on SP active state strategy = self._get_active_parallel_strategy() + paged_adapter = self._active_paged_kv_adapter() + 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( + f"Diffusion paged KV requires an Omni backend with paged support; " + f"selected {self.attn_backend.get_name()}" + ) + if use_paged_attention and strategy is not self._no_parallel_strategy: + strategy_name = strategy.name + if self.use_ring or strategy_name == "ring": + raise NotImplementedError( + "paged Scheduler KV is not supported with Ring attention; use strict Ulysses or no SP" + ) + if strategy_name == "allgather_kv": + raise NotImplementedError("paged Scheduler KV is not supported with AllGather-KV sequence parallelism") + if strategy_name == "ulysses" and get_ulysses_mode(default="strict") != "strict": + raise NotImplementedError("paged Scheduler KV currently supports only strict Ulysses") # 1. Prepare inputs (Communication / Resharding) # For Ulysses: AllToAll Q/K/V; Slicing joint_q/k/v # For Ring: Concat joint_q query, key, value, attn_metadata, ctx = strategy.pre_attention(query, key, value, attn_metadata) - attn_metadata = self._with_kv_cache_dtype(attn_metadata) - - # 2. Kernel Execution (Computation) - if self.use_ring and strategy is not self._no_parallel_strategy: - out = self._run_ring_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. + if use_paged_attention: + assert paged_adapter is not None + paged_kv_context = paged_adapter.prepare_layer_context( + self.prefix, + query, + key, + value, + omni_attn_metadata=attn_metadata, + ) + out = self.attention.forward_paged(paged_kv_context) else: - out = self._run_local_attention(query, key, value, attn_metadata) + attn_metadata = self._with_kv_cache_dtype(attn_metadata) + if self.use_ring and strategy is not self._no_parallel_strategy: + out = self._run_ring_attention(query, key, value, attn_metadata) + else: + out = self._run_local_attention(query, key, value, attn_metadata) # 3. Post-processing (Reverse Communication) # For Ulysses: AllToAll Output, and AllGather Joint Output @@ -331,6 +369,19 @@ class Attention(nn.Module): return out + @staticmethod + def _active_paged_kv_adapter(): + """Return the opaque Worker adapter installed for this forward.""" + + if not is_forward_context_available(): + return None + return getattr(get_forward_context(), "paged_kv_adapter", None) + + def is_paged_kv_active(self) -> bool: + """Return whether this layer will use Scheduler-managed paged KV.""" + + return self.paged_kv_cache_role is not None and self._active_paged_kv_adapter() is not None + def _run_local_attention(self, query, key, value, attn_metadata): self._assert_piecewise_compatible(attn_metadata) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 786e79e9d..5dde3be73 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -714,6 +714,9 @@ class OmniDiffusionConfig: # Local Diffusion KV ownership and cache-layout mode. diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY + # Maximum number of native BlockTable rows one public request can own + # (sequences plus independent contexts). The model adapter defines it. + diffusion_kv_max_rows_per_request: int | None = None # Optional override for the diffusion model runner class (import path). # Precedence in the worker: this override > the runner declared by the @@ -990,6 +993,15 @@ class OmniDiffusionConfig: if not isinstance(self.diffusion_compile_dynamic, bool): raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}") self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode) + if self.diffusion_kv_max_rows_per_request is not None and ( + type(self.diffusion_kv_max_rows_per_request) is not int or self.diffusion_kv_max_rows_per_request <= 0 + ): + raise ValueError("diffusion_kv_max_rows_per_request must be a positive integer when set") + if ( + self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + and self.diffusion_kv_max_rows_per_request is None + ): + raise ValueError("paged_scheduler requires diffusion_kv_max_rows_per_request to be set") if self.kv_cache_memory_bytes is not None and self.kv_cache_memory_bytes < 0: raise ValueError("kv_cache_memory_bytes must be non-negative") if not 0.0 < self.gpu_memory_utilization <= 1.0: diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 4812c8925..d8eb005ae 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -29,7 +29,7 @@ from vllm_omni.diffusion.data import ( DiffusionRequestAbortedError, OmniDiffusionConfig, ) -from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode from vllm_omni.diffusion.diffusion_kv.initialization import initialize_diffusion_kv_control_plane from vllm_omni.diffusion.executor.abstract import DiffusionExecutor from vllm_omni.diffusion.io_support import ( @@ -634,12 +634,25 @@ class DiffusionEngine: if not task.future.done(): task.future.set_exception(exc) + def _remove_diffusion_kv_requests(self, request_ids: Iterable[str]) -> None: + """Clear terminal Worker rows while Scheduler owns the allocations.""" + + od_config = getattr(self, "od_config", None) + if od_config is None or not is_scheduler_paged_kv_mode( + getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + ): + return + unique_request_ids = list(dict.fromkeys(request_ids)) + if unique_request_ids: + self.executor.remove_diffusion_kv_requests(unique_request_ids) + def _emit_finished_outputs( self, finished_ids: set[str], runner_output: BaseRunnerOutput | None = None, missing_result_error: str = "Diffusion execution finished without a final output", ) -> None: + self._remove_diffusion_kv_requests(finished_ids) for rid in finished_ids: if runner_output is not None: _output = runner_output.get_request_output(rid) @@ -659,6 +672,8 @@ class DiffusionEngine: self._emit_finished_outputs(finished_ids, runner_output) return + self._remove_diffusion_kv_requests(finished_ids) + delivered_finished_req_ids: set[str] = set() # finished_ids may have some requests that are not scheduler in this round. @@ -852,6 +867,7 @@ class DiffusionEngine: sched_output = self.scheduler.schedule() if sched_output.is_empty: if target_request_id in sched_output.finished_req_ids: + self._remove_diffusion_kv_requests([target_request_id]) return self._finalize_finished_request(target_request_id) if not self.scheduler.has_requests(): raise RuntimeError("Diffusion scheduler has no runnable requests.") @@ -881,6 +897,7 @@ class DiffusionEngine: if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1: raise ValueError("Sync func should receive one result at one time") if target_request_id in finished_req_ids: + self._remove_diffusion_kv_requests([target_request_id]) req_output = runner_output.get_request_output(target_request_id) output = self._finalize_finished_request( target_request_id, @@ -1232,8 +1249,11 @@ class DiffusionEngine: def _abort_requests(self, request_ids: str | Iterable[str]) -> None: request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids) + request_ids = list(dict.fromkeys(request_ids)) - for request_id in dict.fromkeys(request_ids): + self._remove_diffusion_kv_requests(request_ids) + + for request_id in request_ids: if self.scheduler.get_request_state(request_id) is not None: self.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED) diff --git a/vllm_omni/diffusion/diffusion_kv/config.py b/vllm_omni/diffusion/diffusion_kv/config.py index cd69d8988..0a7d6070b 100644 --- a/vllm_omni/diffusion/diffusion_kv/config.py +++ b/vllm_omni/diffusion/diffusion_kv/config.py @@ -37,3 +37,8 @@ def parse_diffusion_kv_cache_mode(value: object) -> DiffusionKVCacheMode: "use 'dense_legacy' or 'paged_scheduler'" ) return mode + + +def is_scheduler_paged_kv_mode(mode: DiffusionKVCacheMode) -> bool: + """Return whether a parsed cache mode uses Scheduler-owned paging.""" + return mode is DiffusionKVCacheMode.PAGED_SCHEDULER diff --git a/vllm_omni/diffusion/diffusion_kv/initialization.py b/vllm_omni/diffusion/diffusion_kv/initialization.py index 854f3aaa5..31aaaa46e 100644 --- a/vllm_omni/diffusion/diffusion_kv/initialization.py +++ b/vllm_omni/diffusion/diffusion_kv/initialization.py @@ -73,7 +73,13 @@ def initialize_diffusion_kv_control_plane( scheduler_kv_cache_config, vllm_config, ) - executor.set_kv_cache_configs(worker_configs) + # Native sizing may auto-fit ``max_model_len=-1`` against the profiled + # memory budget. Forward that final value so every Worker sizes its + # physical BlockTables from the same bound as the Scheduler. + resolved_max_model_len = vllm_config.model_config.max_model_len + if type(resolved_max_model_len) is not int or resolved_max_model_len <= 0: + raise ValueError("native Diffusion KV sizing must resolve a positive max_model_len") + executor.set_kv_cache_configs(worker_configs, resolved_max_model_len) return ( scheduler_kv_cache_config, scheduler_block_size, diff --git a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py new file mode 100644 index 000000000..479492bac --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py @@ -0,0 +1,647 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import copy +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.distributed import get_dcp_group +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheSpec +from vllm.v1.worker import block_table as native_block_table +from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache +from vllm.v1.worker.gpu.block_table import BlockTables + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode +from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata +from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( + DiffusionPagedAttentionAdapter, + DiffusionPagedAttentionLayerAdapter, + DiffusionPagedAttentionRow, + DiffusionPagedAttentionRowBinding, + PreparedDiffusionPagedAttentionBatch, +) +from vllm_omni.platforms import current_omni_platform + +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 + token_len: int + block_ids: tuple[tuple[int, ...], ...] + + +class DiffusionKVModelRunnerBackend: + """Native paged-KV state and BlockTable operations for a model runner.""" + + def __init__( + self, + *, + vllm_config: VllmConfig, + od_config: OmniDiffusionConfig, + device: torch.device, + ) -> None: + self.vllm_config = vllm_config + self.od_config = od_config + self.device = device + self.kv_cache_config: KVCacheConfig | None = None + self.kv_caches: list[torch.Tensor | list[torch.Tensor]] = [] + self.block_tables: BlockTables | None = None + self.paged_attention_adapter: DiffusionPagedAttentionAdapter | None = None + self._kv_cache_layer_adapters: dict[str, DiffusionPagedAttentionLayerAdapter] = {} + self._diffusion_kv_identity_to_row: dict[DiffusionKVIdentity, int] = {} + self._diffusion_kv_free_rows: list[int] = [] + self._diffusion_kv_request_states: dict[str, _DiffusionKVRequestState] = {} + + def register_kv_cache_layers( + self, + layers: Mapping[str, tuple[Any, KVCacheSpec]], + ) -> dict[str, KVCacheSpec]: + """Register native adapters and return their canonical cache specs.""" + + if self.kv_cache_config is not None: + raise RuntimeError("Diffusion KV cache layers cannot be changed after physical initialization") + + if layers: + current_omni_platform.get_diffusion_kv_block_tables_cls() + if not callable(current_omni_platform.build_diffusion_kv_attn_metadata): + raise NotImplementedError( + "Diffusion paged KV requires platform-native BlockTables and attention metadata hooks" + ) + + forward_context = self.vllm_config.compilation_config.static_forward_context + previous_adapters = self._kv_cache_layer_adapters + for layer_name, (layer, spec) in layers.items(): + if not isinstance(spec, AttentionSpec): + raise TypeError( + f"Diffusion KV layer {layer_name!r} produced unsupported spec {type(spec).__name__}; " + "only native attention specs are supported" + ) + if not layer.attn_backend.supports_paged_kv: + raise NotImplementedError( + f"Diffusion paged KV layer {layer_name!r} requires an Omni backend with paged support; " + f"selected {layer.attn_backend.get_name()}" + ) + existing = forward_context.get(layer_name) + if existing is not None and existing is not layer and existing is not previous_adapters.get(layer_name): + raise RuntimeError(f"Diffusion KV layer name {layer_name!r} collides with the native forward context") + + previous_use_non_causal = self.vllm_config.attention_config.use_non_causal + self.vllm_config.attention_config.use_non_causal = any( + bool(getattr(spec, "non_causal", False)) for _, spec in layers.values() + ) + adapters: dict[str, DiffusionPagedAttentionLayerAdapter] = {} + try: + with set_current_vllm_config(self.vllm_config): + for layer_name, (layer, spec) in layers.items(): + adapters[layer_name] = DiffusionPagedAttentionLayerAdapter( + layer_name=layer_name, + layer=layer, + spec=spec, + vllm_config=self.vllm_config, + device=self.device, + ulysses_degree=self._get_paged_attention_ulysses_degree(layer), + ) + except Exception: + self.vllm_config.attention_config.use_non_causal = previous_use_non_causal + raise + + geometries = {(adapter.num_heads, adapter.num_kv_heads, adapter.head_size) for adapter in adapters.values()} + if len(geometries) > 1: + self.vllm_config.attention_config.use_non_causal = previous_use_non_causal + raise ValueError( + "Native paged attention metadata builders require one rank-local attention geometry; " + f"got {sorted(geometries)!r}" + ) + attention_geometry = next(iter(geometries)) if geometries else None + set_attention_geometry = ( + getattr(self.vllm_config.model_config, "set_attention_geometry", None) + if attention_geometry is not None + else None + ) + if attention_geometry is not None and not callable(set_attention_geometry): + self.vllm_config.attention_config.use_non_causal = previous_use_non_causal + raise RuntimeError( + "Diffusion native paged attention requires a model config that exposes set_attention_geometry" + ) + + updated_forward_context = dict(forward_context) + for layer_name in set(layers) | set(previous_adapters): + entry = updated_forward_context.get(layer_name) + layer_entry = layers.get(layer_name) + if entry is previous_adapters.get(layer_name) or (layer_entry is not None and entry is layer_entry[0]): + updated_forward_context.pop(layer_name, None) + updated_forward_context.update(adapters) + + if attention_geometry is not None: + num_heads, num_kv_heads, head_size = attention_geometry + try: + set_attention_geometry( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_size=head_size, + ) + except Exception: + self.vllm_config.attention_config.use_non_causal = previous_use_non_causal + raise + + forward_context.clear() + forward_context.update(updated_forward_context) + self._kv_cache_layer_adapters = adapters + return {layer_name: adapter.spec for layer_name, adapter in adapters.items()} + + def _get_paged_attention_ulysses_degree(self, layer: Any) -> int: + if bool(getattr(layer, "skip_sequence_parallel", False)): + return 1 + + parallel_config = getattr(self.od_config, "parallel_config", None) + ulysses_degree = int(getattr(parallel_config, "ulysses_degree", 1)) + ring_degree = int(getattr(parallel_config, "ring_degree", 1)) + allgather_degree = int(getattr(parallel_config, "allgather_degree", 1)) + ulysses_mode = str(getattr(parallel_config, "ulysses_mode", "strict")) + if ring_degree > 1: + raise NotImplementedError( + "Diffusion paged attention does not support Ring or hybrid Ulysses+Ring attention" + ) + if allgather_degree > 1: + raise NotImplementedError("Diffusion paged attention does not support AllGather-KV sequence parallelism") + if ulysses_degree > 1 and ulysses_mode != "strict": + raise NotImplementedError( + f"Diffusion paged attention currently supports only strict Ulysses; got ulysses_mode={ulysses_mode!r}" + ) + return ulysses_degree + + def _get_max_rows_per_request(self) -> int: + max_rows = getattr(self.od_config, "diffusion_kv_max_rows_per_request", None) + if type(max_rows) is not int or max_rows <= 0: + raise ValueError( + "paged_scheduler requires a positive integer diffusion_kv_max_rows_per_request adapter limit" + ) + return max_rows + + def _get_max_model_len(self) -> int: + max_len = getattr(self.vllm_config.model_config, "max_model_len", None) + if type(max_len) is not int or max_len <= 0: + raise ValueError("paged_scheduler requires a positive model_config.max_model_len") + return max_len + + def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: + """Allocate and bind native attention KV tensors for this rank.""" + if self.kv_cache_config is not None: + raise RuntimeError("Diffusion native KV cache is already initialized") + if not self._kv_cache_layer_adapters: + raise RuntimeError("Diffusion KV cache layers must be registered before physical initialization") + + configured_layers = { + layer_name for group in kv_cache_config.kv_cache_groups for layer_name in group.layer_names + } + registered_layers = set(self._kv_cache_layer_adapters) + if configured_layers != registered_layers: + raise ValueError( + "Rank-local Diffusion KVCacheConfig layer mismatch: " + f"expected={sorted(registered_layers)}, configured={sorted(configured_layers)}" + ) + + kv_cache_config = copy.deepcopy(kv_cache_config) + max_len = self._get_max_model_len() + max_rows_per_request = self._get_max_rows_per_request() + scheduler_config = self.vllm_config.scheduler_config + max_num_seqs = getattr(scheduler_config, "max_num_seqs", None) + if type(max_num_seqs) is not int or max_num_seqs <= 0: + raise ValueError("paged_scheduler requires a positive scheduler_config.max_num_seqs") + + kv_cache_groups = kv_cache_config.kv_cache_groups + if not kv_cache_groups: + raise ValueError("paged_scheduler requires at least one native KV cache group") + block_sizes = [group.kv_cache_spec.block_size for group in kv_cache_groups] + if any(type(block_size) is not int or block_size <= 0 for block_size in block_sizes): + raise ValueError("native diffusion KV cache groups must have positive integer block sizes") + unaligned_max_num_blocks_per_group = [ + group.kv_cache_spec.max_num_blocks_per_req(self.vllm_config, max_len) for group in kv_cache_groups + ] + if any(type(capacity) is not int or capacity <= 0 for capacity in unaligned_max_num_blocks_per_group): + raise ValueError("native diffusion KV cache groups must have positive row capacities") + max_num_blocks_per_group = [ + native_block_table.get_block_table_width(capacity, block_size) + for capacity, block_size in zip(unaligned_max_num_blocks_per_group, block_sizes, strict=True) + ] + + max_num_reqs = max_num_seqs * max_rows_per_request + max_num_batched_tokens = getattr(scheduler_config, "max_num_batched_tokens", None) + if type(max_num_batched_tokens) is not int or max_num_batched_tokens <= 0: + raise ValueError("scheduler_config.max_num_batched_tokens must be a positive integer") + parallel_config = self.vllm_config.parallel_config + cp_size = parallel_config.decode_context_parallel_size + cp_rank = get_dcp_group().rank_in_group if cp_size > 1 else 0 + cp_interleave = parallel_config.cp_kv_cache_interleave_size + + # Native metadata builders size per-row buffers from max_num_seqs. + # A diffusion public request can occupy several sequence/context rows. + scheduler_config.max_num_seqs = max_num_reqs + kv_caches: list[torch.Tensor | list[torch.Tensor]] = [] + previous_adapter_caches = { + layer_name: adapter.kv_cache for layer_name, adapter in self._kv_cache_layer_adapters.items() + } + try: + with set_current_vllm_config(self.vllm_config): + attn_groups, _, kernel_block_sizes = init_attn_backend( + kv_cache_config, + self.vllm_config, + self.device, + ) + block_tables_cls = current_omni_platform.get_diffusion_kv_block_tables_cls() + block_tables = block_tables_cls( + block_sizes=block_sizes, + max_num_reqs=max_num_reqs, + max_num_batched_tokens=max_num_batched_tokens, + max_num_blocks_per_group=max_num_blocks_per_group, + device=self.device, + kernel_block_sizes=kernel_block_sizes, + cp_size=cp_size, + cp_rank=cp_rank, + cp_interleave=cp_interleave, + ) + paged_attention_adapter = DiffusionPagedAttentionAdapter( + vllm_config=self.vllm_config, + device=self.device, + kv_cache_config=kv_cache_config, + block_tables=block_tables, + attn_groups=attn_groups, + layers=self._kv_cache_layer_adapters, + resolve_row=self._resolve_paged_attention_row, + ) + init_kv_cache( + kv_caches, + self.vllm_config.compilation_config.static_forward_context, + kv_cache_config, + attn_groups, + self.device, + self.vllm_config.cache_config.cache_dtype, + kernel_block_sizes, + self.vllm_config, + ) + except Exception: + scheduler_config.max_num_seqs = max_num_seqs + for layer_name, previous_cache in previous_adapter_caches.items(): + self._kv_cache_layer_adapters[layer_name].kv_cache = previous_cache + raise + + self.kv_caches = kv_caches + self.block_tables = block_tables + self.paged_attention_adapter = paged_attention_adapter + self._diffusion_kv_identity_to_row.clear() + self._diffusion_kv_free_rows = list(range(max_num_reqs - 1, -1, -1)) + self._diffusion_kv_request_states.clear() + self.kv_cache_config = kv_cache_config + + def _validate_row( + self, + *, + identity: DiffusionKVIdentity, + token_len: int, + block_ids: tuple[list[int], ...], + seen_identities: set[DiffusionKVIdentity], + ) -> _DiffusionKVRowInstall: + if identity in seen_identities: + raise ValueError(f"Duplicate diffusion KV row identity: {identity!r}") + seen_identities.add(identity) + if type(token_len) is not int or token_len < 0: + raise ValueError(f"Diffusion KV row {identity!r} has invalid token length {token_len!r}") + + assert self.kv_cache_config is not None + groups = self.kv_cache_config.kv_cache_groups + if not isinstance(block_ids, tuple) or len(block_ids) != len(groups): + actual_groups = len(block_ids) if isinstance(block_ids, (tuple, list)) else "invalid" + raise ValueError(f"Diffusion KV row {identity!r} has {actual_groups} block groups; expected {len(groups)}") + + max_len = self._get_max_model_len() + normalized_groups: list[tuple[int, ...]] = [] + for group_index, (group_block_ids, kv_cache_group) in enumerate(zip(block_ids, groups, strict=True)): + if not isinstance(group_block_ids, list): + raise ValueError(f"Diffusion KV row {identity!r} group {group_index} block IDs must be a list") + spec = kv_cache_group.kv_cache_spec + expected_count = spec.max_num_blocks_per_req(self.vllm_config, token_len) + if len(group_block_ids) != expected_count: + raise ValueError( + f"Diffusion KV row {identity!r} group {group_index} has {len(group_block_ids)} blocks; " + f"expected {expected_count} for {token_len} tokens" + ) + row_capacity = spec.max_num_blocks_per_req(self.vllm_config, max_len) + if len(group_block_ids) > row_capacity: + raise ValueError( + f"Diffusion KV row {identity!r} group {group_index} exceeds row capacity {row_capacity}" + ) + normalized_ids: list[int] = [] + for block_id in group_block_ids: + if type(block_id) is not int: + raise ValueError( + f"Diffusion KV row {identity!r} group {group_index} has non-integer block ID {block_id!r}" + ) + if not 0 <= block_id < self.kv_cache_config.num_blocks: + raise ValueError( + f"Diffusion KV row {identity!r} group {group_index} block ID {block_id} is outside " + f"[0, {self.kv_cache_config.num_blocks})" + ) + normalized_ids.append(block_id) + normalized_groups.append(tuple(normalized_ids)) + + return _DiffusionKVRowInstall( + identity=identity, + token_len=token_len, + block_ids=tuple(normalized_groups), + ) + + def _prepare_install( + self, + metadata: DiffusionKVMetadata, + ) -> tuple[list[_DiffusionKVRowInstall], DiffusionKVSnapshot]: + if self.block_tables is None or self.kv_cache_config is None: + raise RuntimeError("paged_scheduler native KV cache is not initialized") + if type(metadata.request_id) is not str or not metadata.request_id: + raise ValueError("Diffusion KV metadata request_id must be a non-empty string") + if type(metadata.allocation_generation) is not int or metadata.allocation_generation < 0: + raise ValueError("Diffusion KV allocation_generation must be a non-negative integer") + if type(self.kv_cache_config.num_blocks) is not int or self.kv_cache_config.num_blocks <= 0: + raise ValueError("native diffusion KV cache config must provide a positive num_blocks") + if not metadata.sequences: + raise ValueError(f"Diffusion KV request {metadata.request_id!r} must contain at least one sequence") + + seen_identities: set[DiffusionKVIdentity] = set() + installs: list[_DiffusionKVRowInstall] = [] + sequence_snapshots: list[object] = [] + for sequence in metadata.sequences: + if type(sequence.sequence_id) is not int or sequence.sequence_id < 0: + raise ValueError(f"Diffusion KV sequence_id must be a non-negative integer: {sequence.sequence_id!r}") + sequence_identity = (metadata.request_id, sequence.sequence_id, None) + sequence_install = self._validate_row( + identity=sequence_identity, + token_len=sequence.seq_len, + block_ids=sequence.block_ids, + seen_identities=seen_identities, + ) + installs.append(sequence_install) + sequence_snapshots.append( + ( + sequence.sequence_id, + sequence.prefix_len, + sequence.target_len, + sequence.seq_len, + sequence_install.block_ids, + sequence.context_ids, + ) + ) + + context_snapshots: list[object] = [] + for context in metadata.contexts: + if type(context.context_id) is not str or not context.context_id: + raise ValueError("Diffusion KV context_id must be a non-empty string") + context_identity = (metadata.request_id, None, context.context_id) + context_install = self._validate_row( + identity=context_identity, + token_len=context.num_tokens, + block_ids=context.block_ids, + seen_identities=seen_identities, + ) + installs.append(context_install) + context_snapshots.append( + ( + context.context_id, + context.cache_role, + context.num_tokens, + context_install.block_ids, + ) + ) + + max_rows_per_request = self._get_max_rows_per_request() + if len(installs) > max_rows_per_request: + raise ValueError( + f"Diffusion KV request {metadata.request_id!r} requires {len(installs)} rows; " + f"adapter limit is {max_rows_per_request}" + ) + snapshot: DiffusionKVSnapshot = ( + metadata.request_id, + tuple(sequence_snapshots), + tuple(context_snapshots), + ) + return installs, snapshot + + def _rollback_staged_writes(self, rows: list[int], previous_num_blocks: object | None) -> None: + block_tables = self.block_tables + if block_tables is None: + return + for group_table in getattr(block_tables, "block_tables", ()): + clear_staged_writes = getattr(group_table, "clear_staged_writes", None) + if clear_staged_writes is not None: + clear_staged_writes() + if previous_num_blocks is None: + return + native_num_blocks = getattr(block_tables, "num_blocks", None) + if native_num_blocks is None or not hasattr(native_num_blocks, "np"): + return + native_num_blocks.np[:, rows] = previous_num_blocks + copy_to_uva = getattr(native_num_blocks, "copy_to_uva", None) + if copy_to_uva is not None: + copy_to_uva() + + def _apply_rows(self, rows: list[int], installs: list[_DiffusionKVRowInstall]) -> None: + assert self.block_tables is not None + if self.paged_attention_adapter is not None: + self.paged_attention_adapter.invalidate_prepared_batches() + native_num_blocks = getattr(self.block_tables, "num_blocks", None) + previous_num_blocks = None + if native_num_blocks is not None and hasattr(native_num_blocks, "np"): + previous_num_blocks = native_num_blocks.np[:, rows].copy() + try: + for row, install in zip(rows, installs, strict=True): + self.block_tables.append_block_ids( + row, + tuple(list(group_ids) for group_ids in install.block_ids), + overwrite=True, + ) + self.block_tables.apply_staged_writes() + except Exception: + self._rollback_staged_writes(rows, previous_num_blocks) + raise + + def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool: + """Install one Scheduler allocation snapshot into native Worker rows.""" + if not is_scheduler_paged_kv_mode( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + ): + raise ValueError("Dense diffusion execution must not install Diffusion KV metadata") + + installs, snapshot = self._prepare_install(metadata) + current_state = self._diffusion_kv_request_states.get(metadata.request_id) + if current_state is not None: + if metadata.allocation_generation < current_state.generation: + raise ValueError( + f"Stale Diffusion KV allocation generation {metadata.allocation_generation} for request " + f"{metadata.request_id!r}; installed generation is {current_state.generation}" + ) + if snapshot != current_state.snapshot: + raise ValueError( + f"Conflicting Diffusion KV allocation snapshot for active request {metadata.request_id!r}; " + "remove the request before replacing its rows" + ) + if metadata.allocation_generation == current_state.generation: + return False + self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState( + generation=metadata.allocation_generation, + snapshot=snapshot, + row_token_lens=current_state.row_token_lens, + ) + return False + + existing_identities = [ + install.identity for install in installs if install.identity in self._diffusion_kv_identity_to_row + ] + if existing_identities: + raise RuntimeError(f"Diffusion KV row registry contains orphaned identities: {existing_identities!r}") + if len(installs) > len(self._diffusion_kv_free_rows): + raise ValueError( + f"Diffusion KV request {metadata.request_id!r} requires {len(installs)} rows, but only " + f"{len(self._diffusion_kv_free_rows)} rows are free" + ) + + rows = list(reversed(self._diffusion_kv_free_rows[-len(installs) :])) if installs else [] + if rows: + self._apply_rows(rows, installs) + + if installs: + del self._diffusion_kv_free_rows[-len(installs) :] + self._diffusion_kv_identity_to_row.update( + (install.identity, row) for install, row in zip(installs, rows, strict=True) + ) + 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), + ) + return True + + def get_diffusion_kv_row( + self, + request_id: str, + sequence_id: int | None, + context_id: str | None = None, + ) -> int: + """Resolve a Scheduler allocation identity to its native table row.""" + if ( + not is_scheduler_paged_kv_mode( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + ) + or self.block_tables is None + ): + raise RuntimeError("paged_scheduler native BlockTables are not initialized") + identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None) + try: + return self._diffusion_kv_identity_to_row[identity] + except KeyError as exc: + raise KeyError(f"No native diffusion KV row is installed for {identity!r}") from exc + + def get_paged_attention_adapter(self) -> DiffusionPagedAttentionAdapter: + if self.paged_attention_adapter is None: + 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. + + 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) + + def _resolve_paged_attention_row( + self, + request_id: str, + sequence_id: int | None, + context_id: str | None, + ) -> DiffusionPagedAttentionRowBinding: + row_index = self.get_diffusion_kv_row(request_id, sequence_id, context_id) + 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: + return DiffusionPagedAttentionRowBinding( + row_index=row_index, + max_seq_len=token_len, + ) + raise RuntimeError(f"Diffusion KV request state is missing logical length for {identity!r}") + + def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int: + """Retire Worker rows without logically freeing Scheduler-owned blocks.""" + if ( + not is_scheduler_paged_kv_mode( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + ) + or self.block_tables is None + ): + return 0 + request_id_set = set(request_ids) + identities_and_rows = [ + (identity, row) + for identity, row in self._diffusion_kv_identity_to_row.items() + if identity[0] in request_id_set + ] + rows = [row for _, row in identities_and_rows] + if rows: + assert self.kv_cache_config is not None + num_groups = len(self.kv_cache_config.kv_cache_groups) + installs = [ + _DiffusionKVRowInstall( + identity=identity, + token_len=0, + block_ids=tuple(tuple() for _ in range(num_groups)), + ) + for identity, _ in identities_and_rows + ] + self._apply_rows(rows, installs) + + for identity, _ in identities_and_rows: + del self._diffusion_kv_identity_to_row[identity] + for request_id in request_id_set: + self._diffusion_kv_request_states.pop(request_id, None) + self._diffusion_kv_free_rows.extend(rows) + self._diffusion_kv_free_rows.sort(reverse=True) + return len(rows) + + def refresh_block_table_layout(self) -> None: + """Refresh native pointer tensors after a CuMem KV-cache wake-up.""" + if ( + not is_scheduler_paged_kv_mode( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + ) + or self.block_tables is None + ): + return + if self.paged_attention_adapter is not None: + self.paged_attention_adapter.invalidate_prepared_batches() + self.block_tables.init_block_table_layout_tensors() diff --git a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py new file mode 100644 index 000000000..5778af810 --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py @@ -0,0 +1,778 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass, field, replace +from typing import Any + +import torch +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.utils.math_utils import cdiv +from vllm.v1.attention.backend import AttentionType +from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens +from vllm.v1.attention.selector import get_attn_backend +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer +from vllm.v1.worker.gpu.block_table import BlockTables + +from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( + PagedPiecewisePlan, + build_paged_piecewise_plan, +) +from vllm_omni.diffusion.forward_context import override_paged_kv_adapter +from vllm_omni.platforms import current_omni_platform + + +@dataclass(frozen=True) +class DiffusionPagedAttentionRowBinding: + """Worker row and logical length installed for one allocation identity.""" + + row_index: int + max_seq_len: int + + +DiffusionKVRowResolver = Callable[ + [str, int | None, str | None], + DiffusionPagedAttentionRowBinding, +] + + +class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase): + """Register a diffusion layer with vLLM's native cache machinery. + + This object deliberately does *not* subclass ``vllm.Attention``. The + latter owns a second execution path and would bypass Omni's sequence + parallel pre/post hooks. The wrapper only supplies the small + ``AttentionLayerBase`` contract needed by ``init_attn_backend`` and keeps + the platform-native attention implementation/cache view available to the + diffusion adapter. + """ + + def __init__( + self, + *, + layer_name: str, + layer: Any, + spec: AttentionSpec, + vllm_config: VllmConfig, + device: torch.device, + ulysses_degree: int = 1, + ) -> None: + if type(ulysses_degree) is not int or ulysses_degree <= 0: + raise ValueError(f"ulysses_degree must be a positive integer, got {ulysses_degree!r}") + num_heads = int(layer.num_heads) + num_kv_heads = int(spec.num_kv_heads) + if num_heads <= 0 or num_kv_heads <= 0 or num_heads % num_kv_heads != 0: + raise ValueError( + "Paged attention requires positive Q/KV heads with num_heads divisible by num_kv_heads: " + f"num_heads={num_heads}, num_kv_heads={num_kv_heads}" + ) + if num_heads % ulysses_degree != 0 or num_kv_heads % ulysses_degree != 0: + raise ValueError( + "Paged attention requires Q/KV heads divisible by ulysses_degree: " + f"num_heads={num_heads}, num_kv_heads={num_kv_heads}, ulysses_degree={ulysses_degree}" + ) + num_heads //= ulysses_degree + num_kv_heads //= ulysses_degree + + attention_config = vllm_config.attention_config + previous_backend = attention_config.backend + previous_backend_per_kind = attention_config.backend_per_kind + try: + # This is a portable vLLM backend request. The active platform + # resolves it to its native implementation, such as FlashAttention + # on CUDA or AscendAttentionBackend on NPU. + attention_config.backend = AttentionBackendEnum.FLASH_ATTN + attention_config.backend_per_kind = {} + with set_current_vllm_config(vllm_config): + attn_backend = get_attn_backend( + head_size=spec.head_size, + dtype=vllm_config.model_config.dtype, + kv_cache_dtype=vllm_config.cache_config.cache_dtype, + attn_type=AttentionType.DECODER, + num_heads=num_heads, + ) + finally: + attention_config.backend = previous_backend + attention_config.backend_per_kind = previous_backend_per_kind + canonical_spec = replace( + spec, + num_kv_heads=num_kv_heads, + indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(), + ) + self.layer_name = layer_name + self.spec = canonical_spec + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_size = int(canonical_spec.head_size) + self.head_size_v = int(getattr(canonical_spec, "head_size_v", canonical_spec.head_size)) + if self.head_size_v != self.head_size: + raise ValueError( + "Diffusion native paged attention requires equal key/value head sizes; " + f"got head_size={self.head_size}, head_size_v={self.head_size_v}" + ) + self.softmax_scale = float(layer.softmax_scale) + self.attn_backend = attn_backend + self.kv_cache: torch.Tensor | None = None + # Native attention backends consume either device tensors or host + # scalar copies of these scales. The paged diffusion path is currently + # unquantized, so initialize both contracts to identity. + self._q_scale = torch.ones(1, device=device, dtype=torch.float32) + self._k_scale = torch.ones(1, device=device, dtype=torch.float32) + self._v_scale = torch.ones(1, device=device, dtype=torch.float32) + self._q_scale_float = 1.0 + self._k_scale_float = 1.0 + self._v_scale_float = 1.0 + with set_current_vllm_config(vllm_config): + self.impl = self._create_native_impl(vllm_config) + + def get_attn_backend(self): + return self.attn_backend + + def _create_native_impl(self, vllm_config: VllmConfig): + impl_cls = self.attn_backend.get_impl_cls() + impl = impl_cls( + num_heads=self.num_heads, + head_size=self.head_size, + scale=self.softmax_scale, + num_kv_heads=self.num_kv_heads, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype=vllm_config.cache_config.cache_dtype, + logits_soft_cap=None, + attn_type=AttentionType.DECODER, + kv_sharing_target_layer_name=None, + ) + if not callable(getattr(impl, "forward", None)): + 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) + ): + raise RuntimeError( + f"Native paged-attention backend {self.attn_backend.get_name()!r} cannot update the KV cache" + ) + if hasattr(impl, "vllm_flash_attn_version") and impl.vllm_flash_attn_version is None: + raise RuntimeError(f"Native paged-attention kernel is unavailable for diffusion layer {self.layer_name!r}") + return impl + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> AttentionSpec: + del vllm_config + return self.spec + + +@dataclass(frozen=True) +class DiffusionPagedAttentionRow: + """One logical BlockTable row participating in a paged attention call.""" + + request_id: str + query_len: int + seq_len: int + kv_start_pos: int = 0 + sequence_id: int | None = None + context_id: str | None = None + + def __post_init__(self) -> None: + if type(self.request_id) is not str or not self.request_id: + raise ValueError("Paged attention request_id must be a non-empty string") + if (self.sequence_id is None) == (self.context_id is None): + raise ValueError("Paged attention row requires exactly one of sequence_id or context_id") + if self.sequence_id is not None and (type(self.sequence_id) is not int or self.sequence_id < 0): + raise ValueError("Paged attention sequence_id must be a non-negative integer") + if self.context_id is not None and (type(self.context_id) is not str or not self.context_id): + raise ValueError("Paged attention context_id must be a non-empty string") + if type(self.query_len) is not int or self.query_len <= 0: + raise ValueError("Paged attention query_len must be a positive integer") + if type(self.seq_len) is not int or self.seq_len <= 0: + raise ValueError("Paged attention seq_len must be a positive integer") + if type(self.kv_start_pos) is not int or self.kv_start_pos < 0: + raise ValueError("Paged attention kv_start_pos must be a non-negative integer") + if self.kv_start_pos + self.query_len > self.seq_len: + raise ValueError( + "Paged attention write span exceeds seq_len: " + f"start={self.kv_start_pos}, query_len={self.query_len}, seq_len={self.seq_len}" + ) + + @property + def identity(self) -> tuple[str, int | None, str | None]: + return (self.request_id, self.sequence_id, self.context_id) + + +@dataclass(frozen=True) +class PreparedDiffusionPagedAttentionBatch: + """Native metadata shared by all paged attention layers in one forward.""" + + rows: tuple[DiffusionPagedAttentionRow, ...] + row_indices: torch.Tensor + query_start_loc: torch.Tensor + seq_lens: torch.Tensor + positions: torch.Tensor + block_tables: tuple[torch.Tensor, ...] + slot_mappings: torch.Tensor + attn_metadata: dict[str, Any] + slot_mappings_by_layer: dict[str, torch.Tensor] + num_tokens: int + _owner: object = field(repr=False, compare=False) + _generation: int = field(repr=False, compare=False) + + +@dataclass(frozen=True, slots=True) +class DiffusionPagedAttentionContext: + """One layer's native inputs for an Omni paged-backend invocation.""" + + layer: DiffusionPagedAttentionLayerAdapter + query: torch.Tensor + key_write: torch.Tensor + value_write: torch.Tensor + slot_mapping: torch.Tensor + native_metadata: Any + piecewise_plan: PagedPiecewisePlan | None + piecewise_native_metadata: tuple[Any, ...] + query_token_shape: tuple[int, ...] + query_has_head_dims: bool + + def restore_output(self, output: torch.Tensor) -> torch.Tensor: + 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 DiffusionPagedAttentionAdapter: + """Translate diffusion rows into vLLM metadata for an Omni paged backend.""" + + def __init__( + self, + *, + vllm_config: VllmConfig, + device: torch.device, + kv_cache_config: KVCacheConfig, + block_tables: BlockTables, + attn_groups: list[list[Any]], + layers: Mapping[str, DiffusionPagedAttentionLayerAdapter], + resolve_row: DiffusionKVRowResolver, + ) -> None: + if not layers: + raise ValueError("Paged attention requires at least one native attention layer") + if len(attn_groups) != len(kv_cache_config.kv_cache_groups): + raise ValueError( + "Paged attention group mismatch: " + f"builders={len(attn_groups)}, cache_groups={len(kv_cache_config.kv_cache_groups)}" + ) + self.vllm_config = vllm_config + self.device = torch.device(device) + self.kv_cache_config = kv_cache_config + self.block_tables = block_tables + self.attn_groups = attn_groups + self.layers = dict(layers) + self.resolve_row = resolve_row + self._owner = object() + self._prepare_generation = 0 + self._active_batch: PreparedDiffusionPagedAttentionBatch | None = None + self._active_piecewise_plan: PagedPiecewisePlan | None = None + self._active_piecewise_native_metadata: tuple[dict[str, Any], ...] | None = None + self._causal_by_group = self._resolve_group_causality() + self._reorder_batch_threshold = self._resolve_reorder_batch_threshold() + + def _resolve_group_causality(self) -> dict[int, bool]: + causal_by_group: dict[int, bool] = {} + for group_index, cache_group in enumerate(self.kv_cache_config.kv_cache_groups): + layer_causality = { + not bool(getattr(self.layers[layer_name].spec, "non_causal", False)) + for layer_name in cache_group.layer_names + } + if len(layer_causality) != 1: + raise ValueError(f"Paged attention cache group {group_index} mixes causal and non-causal layers") + causal_by_group[group_index] = layer_causality.pop() + return causal_by_group + + def _resolve_reorder_batch_threshold(self) -> int | None: + thresholds = [ + group.get_metadata_builder(0).reorder_batch_threshold + for group_index, groups in enumerate(self.attn_groups) + if self._causal_by_group[group_index] + for group in groups + ] + concrete_thresholds = [threshold for threshold in thresholds if threshold is not None] + return min(concrete_thresholds, default=None) + + def _validate_native_batch_order(self, rows: tuple[DiffusionPagedAttentionRow, ...]) -> None: + threshold = self._reorder_batch_threshold + if threshold is None or not any(self._causal_by_group.values()): + return + + found_long_query = False + for row in rows: + if row.query_len > threshold: + found_long_query = True + elif found_long_query: + raise ValueError( + "Causal paged attention rows must place native decode/short-query rows before " + f"prefill/long-query rows (decode threshold={threshold})" + ) + + def _validate_row_capacity( + self, + rows: tuple[DiffusionPagedAttentionRow, ...], + row_indices: list[int], + ) -> None: + native_num_blocks = self.block_tables.num_blocks.np + blocks_per_kv_block = self.block_tables.blocks_per_kv_block + for row, row_index in zip(rows, row_indices, strict=True): + if type(row_index) is not int or not 0 <= row_index < self.block_tables.max_num_reqs: + raise ValueError(f"Paged attention row {row.identity!r} resolved to invalid Worker row {row_index!r}") + for group_index, (block_size, block_multiplier) in enumerate( + zip( + self.block_tables.block_sizes, + blocks_per_kv_block, + strict=True, + ) + ): + required_blocks = cdiv( + row.seq_len, + block_size * self.block_tables.cp_size, + ) + required_kernel_blocks = required_blocks * block_multiplier + installed_kernel_blocks = int(native_num_blocks[group_index, row_index]) + if required_kernel_blocks > installed_kernel_blocks: + raise ValueError( + f"Paged attention row {row.identity!r} requires {required_blocks} blocks in " + f"cache group {group_index} for seq_len={row.seq_len}, but its installed Worker " + f"row contains only {installed_kernel_blocks // block_multiplier} blocks" + ) + + def _build_native_metadata( + self, + *, + query_lens: Sequence[int], + query_start_loc: torch.Tensor, + query_start_loc_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + positions: torch.Tensor, + block_tables: Sequence[torch.Tensor], + slot_mappings: torch.Tensor, + causal: bool | Mapping[int, bool], + ) -> dict[str, Any]: + dcp_local_seq_lens = None + if self.block_tables.cp_size > 1: + dcp_local_seq_lens = get_dcp_local_seq_lens( + seq_lens, + dcp_size=self.block_tables.cp_size, + dcp_rank=self.block_tables.cp_rank, + cp_kv_cache_interleave_size=self.block_tables.cp_interleave, + ) + with set_current_vllm_config(self.vllm_config): + return current_omni_platform.build_diffusion_kv_attn_metadata( + attn_groups=self.attn_groups, + num_reqs=len(query_lens), + num_tokens=sum(query_lens), + query_start_loc_gpu=query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=max(query_lens), + seq_lens=seq_lens, + max_seq_len=max(int(seq_len) for seq_len in seq_lens_cpu), + block_tables=block_tables, + slot_mappings=slot_mappings, + kv_cache_config=self.kv_cache_config, + seq_lens_cpu=seq_lens_cpu, + seq_lens_cpu_upper_bound=seq_lens_cpu, + dcp_local_seq_lens=dcp_local_seq_lens, + positions=positions, + causal=causal, + ) + + def prepare_batch( + self, + rows: Sequence[DiffusionPagedAttentionRow], + ) -> PreparedDiffusionPagedAttentionBatch: + rows = tuple(rows) + if not rows: + raise ValueError("Paged attention batch must contain at least one row") + if self._active_batch is not None: + raise RuntimeError("Cannot prepare paged attention metadata during an active forward") + if len(rows) > self.block_tables.max_num_reqs: + raise ValueError( + f"Paged attention batch has {len(rows)} rows; Worker capacity is {self.block_tables.max_num_reqs}" + ) + + identities = [row.identity for row in rows] + if len(set(identities)) != len(identities): + raise ValueError("Paged attention batch contains duplicate row identities") + row_bindings = [self.resolve_row(row.request_id, row.sequence_id, row.context_id) for row in rows] + row_indices_list = [binding.row_index for binding in row_bindings] + if len(set(row_indices_list)) != len(row_indices_list): + raise ValueError("Paged attention batch resolves multiple inputs to the same Worker row") + for row, binding in zip(rows, row_bindings, strict=True): + if type(binding.max_seq_len) is not int or binding.max_seq_len < 0: + raise ValueError( + f"Paged attention row {row.identity!r} resolved to invalid logical capacity {binding.max_seq_len!r}" + ) + if row.seq_len > binding.max_seq_len: + raise ValueError( + f"Paged attention row {row.identity!r} uses seq_len={row.seq_len}, but its installed " + f"allocation has logical length {binding.max_seq_len}" + ) + self._validate_native_batch_order(rows) + self._validate_row_capacity(rows, row_indices_list) + + query_lens = [row.query_len for row in rows] + num_tokens = sum(query_lens) + if num_tokens > self.block_tables.max_num_batched_tokens: + raise ValueError( + f"Paged attention batch has {num_tokens} tokens; Worker capacity is " + f"{self.block_tables.max_num_batched_tokens}" + ) + query_offsets = [0] + for query_len in query_lens: + query_offsets.append(query_offsets[-1] + query_len) + + # Native BlockTables reuses persistent gather/slot buffers. Starting a + # new preparation invalidates every batch prepared before it, even if a + # later metadata builder raises. + self._prepare_generation += 1 + generation = self._prepare_generation + row_indices = torch.tensor(row_indices_list, dtype=torch.int32, device=self.device) + query_start_loc_cpu = torch.tensor(query_offsets, dtype=torch.int32) + query_start_loc = query_start_loc_cpu.to(self.device) + seq_lens_cpu = torch.tensor([row.seq_len for row in rows], dtype=torch.int32) + seq_lens = seq_lens_cpu.to(self.device) + positions = torch.cat( + [ + torch.arange( + row.kv_start_pos, + row.kv_start_pos + row.query_len, + dtype=torch.int64, + device=self.device, + ) + for row in rows + ] + ) + + block_tables = self.block_tables.gather_block_tables( + row_indices, + num_reqs_padded=len(rows), + ) + slot_mappings = self.block_tables.compute_slot_mappings( + row_indices, + query_start_loc, + positions, + num_tokens_padded=num_tokens, + ) + causal: bool | Mapping[int, bool] + causal_values = set(self._causal_by_group.values()) + causal = causal_values.pop() if len(causal_values) == 1 else self._causal_by_group + attn_metadata = self._build_native_metadata( + query_lens=query_lens, + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + positions=positions, + block_tables=block_tables, + slot_mappings=slot_mappings, + causal=causal, + ) + slot_mappings_by_layer = build_slot_mappings_by_layer( + slot_mappings, + self.kv_cache_config, + ) + return PreparedDiffusionPagedAttentionBatch( + rows=rows, + row_indices=row_indices, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + positions=positions, + block_tables=tuple(block_tables), + slot_mappings=slot_mappings, + attn_metadata=attn_metadata, + slot_mappings_by_layer=slot_mappings_by_layer, + num_tokens=num_tokens, + _owner=self._owner, + _generation=generation, + ) + + def invalidate_prepared_batches(self) -> None: + """Invalidate native buffer views after BlockTable state changes.""" + + if self._active_batch is not None: + raise RuntimeError("Cannot change paged attention BlockTables during an active forward") + self._prepare_generation += 1 + + @contextmanager + def activate( + self, + batch: PreparedDiffusionPagedAttentionBatch, + ) -> Iterator[DiffusionPagedAttentionAdapter]: + if batch._owner is not self._owner: + raise ValueError("Prepared paged attention batch belongs to a different adapter") + if batch._generation != self._prepare_generation: + raise ValueError("Prepared paged attention batch is stale after a newer batch preparation") + if self._active_batch is not None: + raise RuntimeError("Paged attention adapter already has an active forward") + self._active_batch = batch + self._active_piecewise_plan = None + self._active_piecewise_native_metadata = None + try: + with override_paged_kv_adapter(self), set_current_vllm_config(self.vllm_config): + yield self + finally: + self._active_piecewise_plan = None + self._active_piecewise_native_metadata = None + self._active_batch = None + + @staticmethod + def _flatten_tensor( + tensor: torch.Tensor, + *, + num_heads: int, + head_size: int, + name: str, + ) -> tuple[torch.Tensor, tuple[int, ...], bool]: + if tensor.ndim >= 3 and tensor.shape[-2:] == (num_heads, head_size): + return tensor.reshape(-1, num_heads, head_size), tuple(tensor.shape[:-2]), True + hidden_size = num_heads * head_size + if tensor.ndim >= 1 and tensor.shape[-1] == hidden_size: + return tensor.reshape(-1, num_heads, head_size), tuple(tensor.shape[:-1]), False + raise ValueError( + f"Paged attention {name} must end in ({num_heads}, {head_size}) or ({hidden_size},); " + f"got shape={tuple(tensor.shape)}" + ) + + @staticmethod + def _validate_token_layout( + token_shape: tuple[int, ...], + batch: PreparedDiffusionPagedAttentionBatch, + *, + name: str, + ) -> None: + if len(token_shape) == 1: + if token_shape[0] == batch.num_tokens: + return + raise ValueError( + f"Paged attention {name} token count must match the prepared write batch: " + f"tokens={token_shape[0]}, prepared={batch.num_tokens}" + ) + 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}" + ) + raise ValueError( + f"Paged attention {name} supports packed [T, ...] or uniform batched [B, T, ...] token layouts; " + f"got token shape={token_shape}" + ) + + def _get_piecewise_plan( + self, + full_attn_spans: list[list[tuple[int, int]]], + ) -> PagedPiecewisePlan: + batch = self._active_batch + if batch is None: + raise RuntimeError("Piecewise paged attention requires an active prepared batch") + if any(row.kv_start_pos + row.query_len != row.seq_len for row in batch.rows): + invalid_rows = [row.identity for row in batch.rows if row.kv_start_pos + row.query_len != row.seq_len] + raise ValueError( + "Paged piecewise attention requires each query/write span to end at seq_len; " + f"invalid rows={invalid_rows!r}" + ) + + cached_plan = self._active_piecewise_plan + if cached_plan is not None: + frozen_spans = tuple(tuple(tuple(span) for span in row_spans) for row_spans in full_attn_spans) + if cached_plan.spans != frozen_spans: + raise ValueError("Paged piecewise attention metadata changed between layers in one active batch") + return cached_plan + + plan = build_paged_piecewise_plan( + full_attn_spans, + query_offsets=[row.kv_start_pos for row in batch.rows], + query_lens=[row.query_len for row in batch.rows], + seq_lens=[row.seq_len for row in batch.rows], + device=self.device, + ) + self._active_piecewise_plan = plan + return plan + + def _get_piecewise_native_metadata( + self, + plan: PagedPiecewisePlan, + ) -> tuple[dict[str, Any], ...]: + """Build one native metadata object per segment and cache it by batch.""" + + batch = self._active_batch + if batch is None: + raise RuntimeError("Piecewise paged attention requires an active prepared batch") + cached_metadata = self._active_piecewise_native_metadata + if cached_metadata is not None: + return cached_metadata + + metadata_by_segment: list[dict[str, Any]] = [] + for packed_segment in plan.segments: + row_segments = packed_segment.row_segments + query_lens = [segment.q_end - segment.q_start for segment in row_segments] + seq_lens_values = [segment.kv_end for segment in row_segments] + query_offsets = [0] + for query_len in query_lens: + 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"), + ) + ) + self._active_piecewise_native_metadata = tuple(metadata_by_segment) + return self._active_piecewise_native_metadata + + def prepare_layer_context( + self, + layer_name: str, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + omni_attn_metadata: Any | None = None, + ) -> DiffusionPagedAttentionContext: + if self._active_batch is None: + raise RuntimeError("Paged attention forward must run inside adapter.activate(batch)") + try: + layer = self.layers[layer_name] + except KeyError as exc: + raise KeyError(f"Unknown diffusion paged attention layer {layer_name!r}") from exc + batch = self._active_batch + full_attn_spans = self._validate_omni_attn_metadata(omni_attn_metadata) + + query_flat, query_token_shape, query_has_head_dims = self._flatten_tensor( + query, + num_heads=layer.num_heads, + head_size=layer.head_size, + name="query", + ) + key_flat, key_token_shape, _ = self._flatten_tensor( + key, + num_heads=layer.num_kv_heads, + head_size=layer.head_size, + name="key", + ) + value_flat, value_token_shape, _ = self._flatten_tensor( + value, + num_heads=layer.num_kv_heads, + 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 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}; " + f"got query={query.device}, key={key.device}, value={value.device}" + ) + if query.dtype != key.dtype or query.dtype != value.dtype: + raise ValueError(f"Paged attention Q/K/V dtypes must match; got {query.dtype}, {key.dtype}, {value.dtype}") + expected_dtype = self.vllm_config.model_config.dtype + if query.dtype != expected_dtype: + raise ValueError( + f"Paged attention Q/K/V dtype must match model activation dtype {expected_dtype}; got {query.dtype}" + ) + if not bool(getattr(layer.spec, "non_causal", False)): + non_suffix_rows = [row.identity for row in batch.rows if row.kv_start_pos + row.query_len != row.seq_len] + if non_suffix_rows: + raise ValueError( + "Causal paged attention requires each query/write span to end at seq_len; " + f"invalid rows={non_suffix_rows!r}" + ) + + slot_mapping = batch.slot_mappings_by_layer.get(layer_name) + if slot_mapping is None: + raise KeyError(f"No native slot mapping was built for diffusion layer {layer_name!r}") + slot_mapping = slot_mapping.reshape(-1)[: batch.num_tokens] + if slot_mapping.numel() != batch.num_tokens: + raise ValueError( + f"Paged attention slot mapping for {layer_name!r} has {slot_mapping.numel()} entries; " + f"expected {batch.num_tokens}" + ) + + # vLLM's writer 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: + raise KeyError(f"No native attention metadata was built for diffusion layer {layer_name!r}") from exc + piecewise_plan: PagedPiecewisePlan | None = None + piecewise_native_metadata: tuple[Any, ...] = () + if full_attn_spans is not None: + piecewise_plan = self._get_piecewise_plan(full_attn_spans) + native_metadata_by_segment = self._get_piecewise_native_metadata(piecewise_plan) + try: + piecewise_native_metadata = tuple( + segment_metadata[layer_name] for segment_metadata in native_metadata_by_segment + ) + except KeyError as exc: + raise KeyError( + f"No piecewise native attention metadata was built for diffusion layer {layer_name!r}" + ) from exc + return DiffusionPagedAttentionContext( + layer=layer, + query=query_flat.contiguous(), + key_write=key_flat, + value_write=value_flat, + slot_mapping=slot_mapping, + native_metadata=native_metadata, + piecewise_plan=piecewise_plan, + piecewise_native_metadata=piecewise_native_metadata, + query_token_shape=query_token_shape, + query_has_head_dims=query_has_head_dims, + ) + + @staticmethod + def _validate_omni_attn_metadata( + metadata: Any | None, + ) -> list[list[tuple[int, int]]] | None: + if metadata is None: + return None + full_attn_spans = getattr(metadata, "full_attn_spans", None) + attn_mask = getattr(metadata, "attn_mask", None) + unsupported_fields = [ + field_name + for field_name in ( + "joint_attn_mask", + "query_ranges", + "video_layout", + ) + if getattr(metadata, field_name, None) is not None + ] + if attn_mask is not None: + if full_attn_spans is None: + unsupported_fields.append("attn_mask") + elif not isinstance(attn_mask, torch.Tensor) or attn_mask.ndim != 4: + unsupported_fields.append("attn_mask (piecewise paging requires a 4D tensor)") + extra = getattr(metadata, "extra", None) + if extra: + unsupported_fields.append(f"extra={sorted(extra)}") + if unsupported_fields: + raise NotImplementedError( + "Diffusion paged attention cannot translate Omni attention metadata fields " + f"{unsupported_fields!r} to native FlashAttention metadata" + ) + return full_attn_spans diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index e653dc7ce..a491c3391 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/vllm_omni/diffusion/executor/abstract.py @@ -142,13 +142,24 @@ class DiffusionExecutor(ABC): raise TypeError("determine_available_kv_memory must return list[int]") return result - def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig]) -> None: - """Send each native rank-local config back to all Workers.""" + def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig], resolved_max_model_len: int) -> None: + """Send rank-local configs and the resolved model length to all Workers.""" # The default control-plane RPC mode executes on every rank and has # rank 0 return the gathered rank statuses, so failures on nonzero # ranks are not silently dropped. - self.collective_rpc("set_kv_cache_configs", args=(kv_cache_configs,)) + self.collective_rpc("set_kv_cache_configs", args=(kv_cache_configs, resolved_max_model_len)) + + def remove_diffusion_kv_requests(self, request_ids: list[str]) -> None: + """Clear request rows on every Worker after Scheduler retirement.""" + + unique_request_ids = list(dict.fromkeys(request_ids)) + if not unique_request_ids: + return + self.collective_rpc( + "remove_diffusion_kv_requests", + args=(unique_request_ids,), + ) @abstractmethod def shutdown(self) -> None: diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index bb589d7bb..af578d662 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -2,7 +2,7 @@ from __future__ import annotations from contextlib import contextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch import vllm.ir @@ -26,6 +26,11 @@ class ForwardContext: vllm_config: VllmConfig | None = None omni_diffusion_config: OmniDiffusionConfig | None = None attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | 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 split_text_embed_in_sp: bool = False denoise_step_idx: int | None = None denoise_timestep: float | None = None @@ -209,6 +214,30 @@ def set_forward_context( yield +@contextmanager +def override_paged_kv_adapter(adapter: Any | None): + """Temporarily expose a Worker paged-KV adapter to Omni Attention. + + This is deliberately a small context override instead of a second global + forward context. The model runner owns the outer context and the adapter + only replaces one opaque field while its prepared native metadata is live. + """ + + if _forward_context is None: + # Unit-level adapter users can prepare/activate metadata without an + # Omni model forward context. In that case there is nothing to + # override and the adapter's explicit ``forward`` API remains usable. + yield + return + + previous = _forward_context.paged_kv_adapter + _forward_context.paged_kv_adapter = adapter + try: + yield + finally: + _forward_context.paged_kv_adapter = previous + + 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: 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 9328abe01..2cc42df71 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -1134,7 +1134,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: + if not keep_kv_compressed and not self.attn.is_paged_kv_active(): key = repeat_kv(key, repeat_num) value = repeat_kv(value, repeat_num) if self.sp_size > 1: diff --git a/vllm_omni/diffusion/vllm_config.py b/vllm_omni/diffusion/vllm_config.py index b56109de1..85e288422 100644 --- a/vllm_omni/diffusion/vllm_config.py +++ b/vllm_omni/diffusion/vllm_config.py @@ -5,7 +5,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any import torch @@ -80,6 +80,15 @@ class _DiffusionVllmModelConfig: logits_processors: tuple[str, ...] = () use_mla: bool = False is_moe: bool = False + # Native vLLM attention layers/builders read these fields directly. The + # diffusion model does not have a full vLLM ModelConfig, so keep the + # narrow compatibility surface here instead of manufacturing a second + # model configuration object in the Worker backend. + is_mm_prefix_lm: bool = False + rswa_window: int | None = None + _attention_num_heads: int | None = field(default=None, init=False, repr=False) + _attention_num_kv_heads: int | None = field(default=None, init=False, repr=False) + _attention_head_size: int | None = field(default=None, init=False, repr=False) # Needed for models that bundle things like LogitsProcessors, e.g., SenseNova head_dtype: torch.dtype | None = None @@ -91,6 +100,56 @@ class _DiffusionVllmModelConfig: def is_nvfp4_quantized(self) -> bool: return self.quantization == "modelopt_fp4" + def set_attention_geometry( + self, + *, + num_heads: int, + num_kv_heads: int, + head_size: int, + ) -> None: + """Record rank-local geometry needed by native metadata builders. + + Diffusion attention modules are loaded outside vLLM's model loader, so + their dimensions are not available when this lightweight config is + created. The paged Worker backend fills them in after model loading + and before calling ``init_attn_backend``. + """ + + values = { + "num_heads": num_heads, + "num_kv_heads": num_kv_heads, + "head_size": head_size, + } + invalid = {name: value for name, value in values.items() if type(value) is not int or value <= 0} + if invalid: + raise ValueError(f"Diffusion attention geometry must contain positive integers: {invalid!r}") + self._attention_num_heads = num_heads + self._attention_num_kv_heads = num_kv_heads + self._attention_head_size = head_size + + def _require_attention_geometry(self) -> tuple[int, int, int]: + geometry = ( + self._attention_num_heads, + self._attention_num_kv_heads, + self._attention_head_size, + ) + if any(value is None for value in geometry): + raise RuntimeError("Diffusion attention geometry is unavailable before paged cache layer registration") + num_heads, num_kv_heads, head_size = geometry + assert num_heads is not None and num_kv_heads is not None and head_size is not None + return num_heads, num_kv_heads, head_size + + def get_num_attention_heads(self, parallel_config: Any) -> int: + del parallel_config + return self._require_attention_geometry()[0] + + def get_num_kv_heads(self, parallel_config: Any) -> int: + del parallel_config + return self._require_attention_geometry()[1] + + def get_head_size(self) -> int: + return self._require_attention_geometry()[2] + @property def is_diffusion(self) -> bool: # This config is also installed as the current VllmConfig for the @@ -164,6 +223,7 @@ def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiff vllm_config.cache_config.kv_cache_memory_bytes = getattr(od_config, "kv_cache_memory_bytes", None) # Prefix hashes/publication are intentionally deferred to the next PR. vllm_config.cache_config.enable_prefix_caching = False + vllm_config.scheduler_config.max_num_seqs = int(od_config.max_num_seqs) 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) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 44993fd2c..2992d2bde 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -24,6 +24,7 @@ from vllm.logger import init_logger from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec +from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import cache_summary from vllm_omni.diffusion.cache.prompt_embed_cache import ( install_prompt_embed_cache, @@ -34,6 +35,11 @@ from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig 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 ( + DiffusionPagedAttentionRow, + PreparedDiffusionPagedAttentionBatch, +) 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 ( @@ -144,9 +150,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): self.prompt_embed_cache: Any | None = None self.input_batch: InputBatch | None = None self.model_memory_usage = 0 - # Rank-local native config generated by the Engine. Physical tensors - # and BlockTable binding remain the responsibility of the paged - # data-plane initialization path. + self.diffusion_kv_backend = DiffusionKVModelRunnerBackend( + vllm_config=vllm_config, + od_config=od_config, + device=device, + ) + # Compatibility view for callers that inspect the installed config; + # physical ownership remains in ``diffusion_kv_backend``. self.kv_cache_config: KVCacheConfig | None = None # Cache for per-request stepwise state. @@ -381,41 +391,60 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: """Collect native specs from cache-enabled loaded attention modules.""" - from vllm_omni.diffusion.attention.layer import Attention - from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode - if self.od_config.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER: return {} if self.pipeline is None: raise RuntimeError("Model must be loaded before collecting Diffusion KV cache specs") - specs: dict[str, KVCacheSpec] = {} + cache_layers: dict[str, tuple[Attention, KVCacheSpec]] = {} for layer_name, 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: - specs[layer_name] = spec - if not specs: + cache_layers[layer_name] = (module, spec) + if not cache_layers: raise RuntimeError( "paged_scheduler Diffusion KV found no cache-enabled Attention modules " f"in {type(self.pipeline).__name__}" ) - return specs + return self.diffusion_kv_backend.register_kv_cache_layers(cache_layers) def set_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None: - """Retain the Engine-generated rank-local config for the data plane.""" + """Physically initialize the Engine-generated rank-local config.""" - configured_layers = { - layer_name for group in kv_cache_config.kv_cache_groups for layer_name in group.layer_names - } - expected_layers = set(self.get_kv_cache_spec()) - if configured_layers != expected_layers: - raise ValueError( - "Rank-local Diffusion KVCacheConfig layer mismatch: " - f"expected={sorted(expected_layers)}, configured={sorted(configured_layers)}" - ) - self.kv_cache_config = kv_cache_config + self.diffusion_kv_backend.initialize_kv_cache(kv_cache_config) + self.kv_cache_config = self.diffusion_kv_backend.kv_cache_config + + def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool: + return self.diffusion_kv_backend.install_diffusion_kv_metadata(metadata) + + def get_diffusion_kv_row( + self, + request_id: str, + sequence_id: int | None, + context_id: str | None = None, + ) -> int: + return self.diffusion_kv_backend.get_diffusion_kv_row(request_id, sequence_id, context_id) + + def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int: + return self.diffusion_kv_backend.remove_diffusion_kv_requests(request_ids) + + def refresh_diffusion_kv_block_table_layout(self) -> None: + self.diffusion_kv_backend.refresh_block_table_layout() + + def prepare_paged_attention_batch( + self, + rows: list[DiffusionPagedAttentionRow] | tuple[DiffusionPagedAttentionRow, ...], + ) -> PreparedDiffusionPagedAttentionBatch: + """Prepare model-specific query/write spans against installed Worker 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) def clear_prompt_embed_cache(self) -> None: """Evict all cached text-encoder outputs (e.g. between training epochs). @@ -664,17 +693,25 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): request_id=req.request_id, metadata=diffusion_kv_metadata, ) - runner_output = self._execute_request_list( - [req], - od_config=self.od_config, - allow_single_output=True, - require_request_batch_support=False, - kv_prefetch_job=kv_prefetch_job, - record_name="pipeline_forward", - ) - output = runner_output.runner_outputs[0].result - assert output is not None - return output + installed_request = False + if diffusion_kv_metadata is not None: + self.install_diffusion_kv_metadata(diffusion_kv_metadata) + installed_request = True + try: + runner_output = self._execute_request_list( + [req], + od_config=self.od_config, + allow_single_output=True, + require_request_batch_support=False, + kv_prefetch_job=kv_prefetch_job, + record_name="pipeline_forward", + ) + output = runner_output.runner_outputs[0].result + assert output is not None + return output + finally: + if installed_request: + self.remove_diffusion_kv_requests([req.request_id]) def profile_run(self, requests: list[OmniDiffusionRequest]) -> None: """Run the maximum per-rank request batch for memory profiling. @@ -746,14 +783,23 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): request_id=new_req.req.request_id, metadata=new_req.diffusion_kv_metadata, ) - reqs = [nr.req for nr in scheduler_output.scheduled_new_reqs] - return self._execute_request_list( - reqs, - od_config=od_config, - allow_single_output=False, - require_request_batch_support=True, - record_name="pipeline_forward_batch", - ) + installed_request_ids: list[str] = [] + try: + for new_req in scheduler_output.scheduled_new_reqs: + if new_req.diffusion_kv_metadata is not None: + self.install_diffusion_kv_metadata(new_req.diffusion_kv_metadata) + installed_request_ids.append(new_req.request_id) + reqs = [nr.req for nr in scheduler_output.scheduled_new_reqs] + return self._execute_request_list( + reqs, + od_config=od_config, + allow_single_output=False, + require_request_batch_support=True, + record_name="pipeline_forward_batch", + ) + finally: + if installed_request_ids: + self.remove_diffusion_kv_requests(installed_request_ids) # ------------------------------------------------------------------ # Step-wise execution @@ -884,6 +930,40 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): if self.od_config.cache_backend not in (None, "none"): raise ValueError("Step mode does not support cache_backend yet.") + # Scheduler metadata is installed only for newly admitted requests; + # cached requests continue to use the row installed on their first + # step. Paged rows are retired here for scheduler-finished requests so + # a subsequent wave can reuse the native table slots. Dense execution + # must not issue Worker KV cleanup side effects. + if ( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + and scheduler_output.finished_req_ids + ): + self.remove_diffusion_kv_requests(list(scheduler_output.finished_req_ids)) + installed_request_ids: list[str] = [] + try: + for new_req in scheduler_output.scheduled_new_reqs: + if new_req.diffusion_kv_metadata is not None: + self.install_diffusion_kv_metadata(new_req.diffusion_kv_metadata) + installed_request_ids.append(new_req.request_id) + return self._execute_stepwise_core( + scheduler_output, + record_output_peak_memory=record_output_peak_memory, + ) + except Exception: + if installed_request_ids: + self.remove_diffusion_kv_requests(installed_request_ids) + raise + + def _execute_stepwise_core( + self, + scheduler_output: DiffusionSchedulerOutput, + *, + record_output_peak_memory: bool, + ) -> BatchRunnerOutput: + """Run the denoise step after metadata admission and row installation.""" + use_hsdp = self.od_config.parallel_config.use_hsdp grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() with grad_context: @@ -1005,7 +1085,16 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): state.peak_memory_mb, ) + terminal_request_ids = [ + runner_output.request_id for runner_output in runner_output_list if runner_output.finished + ] self._update_states_after(states, input_batch, pipeline_interrupted) + if ( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + and terminal_request_ids + ): + self.remove_diffusion_kv_requests(terminal_request_ids) return BatchRunnerOutput.from_list(runner_output_list) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index faa35002d..9344cb9b4 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -188,6 +188,11 @@ 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), @@ -431,16 +436,38 @@ class DiffusionWorker: ) ] - def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig]) -> None: - """Install this rank's native config without allocating tensors yet.""" + def set_kv_cache_configs( + self, + kv_cache_configs: list[KVCacheConfig], + resolved_max_model_len: int, + ) -> None: + """Select this rank's config and initialize its physical KV pages.""" assert self.model_runner is not None + assert self.vllm_config is not None if len(kv_cache_configs) != self.od_config.num_gpus: raise ValueError( "Diffusion KVCacheConfig rank count mismatch: " f"expected={self.od_config.num_gpus}, got={len(kv_cache_configs)}" ) - self.model_runner.set_kv_cache_config(kv_cache_configs[self.rank]) + if not 0 <= self.rank < len(kv_cache_configs): + raise ValueError(f"Diffusion Worker rank {self.rank} has no rank-local KVCacheConfig") + if type(resolved_max_model_len) is not int or resolved_max_model_len <= 0: + raise ValueError("resolved Diffusion KV max_model_len must be a positive integer") + + # Native cache sizing may resolve an explicit ``-1`` model-length + # sentinel to the capacity that actually fits the profiled pool. + self.vllm_config.model_config.max_model_len = resolved_max_model_len + kv_cache_config = kv_cache_configs[self.rank] + self.vllm_config.cache_config.num_gpu_blocks = kv_cache_config.num_blocks + with self._maybe_get_memory_pool_context("kv_cache"): + self.model_runner.set_kv_cache_config(kv_cache_config) + + def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int: + """Clear Worker-local rows without freeing Scheduler-owned blocks.""" + + assert self.model_runner is not None, "Model runner not initialized" + return self.model_runner.remove_diffusion_kv_requests(request_ids) def init_lora_manager(self) -> None: """Initialize the LoRA manager for this worker.""" @@ -730,6 +757,8 @@ class DiffusionWorker: allocator = CuMemAllocator.get_instance() allocator.wake_up(tags) current_omni_platform.synchronize() + if self.model_runner is not None and (tags is None or "kv_cache" in tags): + self.model_runner.refresh_diffusion_kv_block_table_layout() if len(self._sleep_saved_buffers) and self.model_runner is not None: model = self.model_runner.pipeline for name, buffer in model.named_buffers(): diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 7759cb6d1..30421fa58 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -184,6 +184,28 @@ class OmniPlatform(Platform): """ return "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner" + @classmethod + def get_diffusion_kv_block_tables_cls(cls) -> type: + """Return the platform's native paged-KV BlockTables implementation.""" + from vllm.v1.worker.gpu.block_table import BlockTables + + return BlockTables + + @classmethod + 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. + """ + from vllm.v1.worker.gpu.attn_utils import build_attn_metadata + + kwargs.pop("seq_lens_cpu", None) + return build_attn_metadata(**kwargs) + @classmethod def init_diffusion_worker_vllm_config( cls, diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 9a6a278e4..8ad695fed 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -80,6 +80,27 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): init_ascend_config(vllm_config) + @classmethod + def get_diffusion_kv_block_tables_cls(cls) -> type: + from vllm_ascend.worker.v2.block_table import AscendBlockTables + + return AscendBlockTables + + @classmethod + def build_diffusion_kv_attn_metadata(cls, **kwargs: Any) -> dict[str, Any]: + """Build the Ascend metadata required by the native NPU backend.""" + from vllm_ascend.attention.attention_v1 import AscendAttentionState + from vllm_ascend.worker.v2.attn_utils import build_attn_metadata + + kwargs = dict(kwargs) + seq_lens_cpu = kwargs.pop("seq_lens_cpu") + kwargs["seq_lens_np"] = seq_lens_cpu.detach().cpu().numpy() + # The diffusion adapter always supplies a paged cache and the current + # K/V write span. ChunkedPrefill is Ascend's cache-backed FIA state for + # both multi-token updates and single-token updates in this path. + kwargs["attn_state"] = AscendAttentionState.ChunkedPrefill + return build_attn_metadata(**kwargs) + @classmethod def init_diffusion_model_runner_runtime(cls, vllm_config: Any, od_config: Any, device: torch.device) -> None: from vllm_ascend.ascend_forward_context import set_mc2_mask, set_mc2_tokens_capacity