# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import weakref import pytest import torch from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.diffusion.models.qwen_image_21 import decode_graph from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] class DecodeModel(torch.nn.Module): in_channels = 4 def _decode_graph_forward(self, entry): return ( entry.hidden.sin() + entry.k[0].mean(dim=(1, 2, 3))[:, None, None] + entry.v[0].mean(dim=(1, 2, 3))[:, None, None] + entry.freqs[None, :, :] ) @pytest.fixture def graph_case(request, monkeypatch): device = "cuda" if request.node.get_closest_marker("cuda") else "cpu" manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), max_entries=2) monkeypatch.setattr(manager, "eligible", lambda: True) monkeypatch.setattr(manager, "_backend_name", lambda: "test") cleanup = [] monkeypatch.setattr(decode_graph.current_omni_platform, "empty_cache", lambda: cleanup.append(True)) return manager, device, cleanup def make_cache(device, value): return [ { "cond": { "key": torch.full((1, 2, 1, 4), value, device=device), "value": torch.full((1, 2, 1, 4), value + 1, device=device), } } ] def register(manager, cache, layout=(1, 2, 2)): device = cache[0]["cond"]["key"].device manager.register_prefill( kv_cache=cache, cache_branch="cond", prefix_len=2, img_shapes=[layout], target_freqs=torch.full((4, 4), float(layout[1]), device=device), joint_key_valid=None, dtype=torch.float32, ) def decode(manager, cache, layout=(1, 2, 2)): device = cache[0]["cond"]["key"].device return manager.try_decode( hidden_states=torch.ones(1, 4, 4, device=device), timestep=torch.ones(1, device=device), kv_cache=cache, cache_branch="cond", img_shapes=[layout], img_mask=torch.tensor([[False, False, True]], device=device), encoder_hidden_states_mask=None, ) @pytest.mark.cpu def test_model_level_offload_without_staging_disables_decode_graphs(): manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), model_level_offload=True) assert manager._offload_reason() is not None assert manager.eligible() is False @pytest.mark.cpu def test_top_level_offload_hook_without_staging_disables_decode_graphs(): model = DecodeModel() manager = decode_graph.QwenImage21DecodeGraphManager(model) registry = HookRegistry.get_or_create(model) # A present-but-empty registry (e.g. after hooks were removed) must not disqualify. assert manager._offload_reason() is None registry.register_hook("sequential_offload", ModelHook()) assert manager._offload_reason() is not None assert manager.eligible() is False @pytest.mark.cpu def test_persistent_staging_makes_model_level_offload_graph_eligible(): model = DecodeModel() manager = decode_graph.QwenImage21DecodeGraphManager(model, model_level_offload=True) hook = SequentialOffloadHook(offload_targets=[], device=torch.device("cpu")) HookRegistry.get_or_create(model).register_hook(SequentialOffloadHook._HOOK_NAME, hook) # No staging storage yet: a swap would allocate fresh storage per generation. assert manager._offload_reason() is not None # Staging established: fixed device storage keeps captured pointers valid. hook._stager = object() assert manager._offload_reason() is None # A second, non-offload hook on the top-level module still disqualifies. model._hook_registry.register_hook("teacache", ModelHook()) assert manager._offload_reason() is not None @pytest.mark.cpu def test_allocation_oom_releases_partial_buffers_and_remembers_failure(graph_case, monkeypatch): manager, device, cleanup = graph_case cache = make_cache(device, 1.0) constructor = decode_graph.QwenImage21DecodeGraphEntry allocated = [] def fail_allocation(**kwargs): entry = constructor(**kwargs) allocated.append(weakref.ref(entry.k[0])) raise torch.OutOfMemoryError("injected allocation OOM") monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", fail_allocation) register(manager, cache) register(manager, cache) assert len(allocated) == 1 assert allocated[0]() is None assert list(manager.entries.values()) == [None] assert len(cleanup) == 1 assert torch.equal(cache[0]["cond"]["key"], torch.ones(1, 2, 1, 4)) @pytest.mark.cpu def test_evict_before_allocating_replacement(graph_case, monkeypatch): manager, device, _ = graph_case manager.max_entries = 1 cache = make_cache(device, 1.0) register(manager, cache) old_buffer = weakref.ref(next(iter(manager.entries.values())).k[0]) constructor = decode_graph.QwenImage21DecodeGraphEntry def allocate_after_release(**kwargs): assert old_buffer() is None return constructor(**kwargs) monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", allocate_after_release) register(manager, cache, (1, 1, 4)) assert len(manager.entries) == 1 @pytest.mark.cpu def test_registration_preserves_request_cache_and_separates_layouts(graph_case): manager, device, _ = graph_case first, second = make_cache(device, 1.0), make_cache(device, 5.0) original = first[0]["cond"]["key"] register(manager, first) register(manager, second) assert first[0]["cond"]["key"] is original assert torch.equal(original, torch.ones_like(original)) register(manager, second, (1, 1, 4)) assert len(manager.entries) == 2 entries = list(manager.entries.values()) assert not torch.equal(entries[0].freqs, entries[1].freqs) @pytest.mark.cuda @pytest.mark.gpu @pytest.mark.parametrize("failure_phase", ["warmup", "capture"]) def test_capture_failure_releases_buffers_and_keeps_request_cache(graph_case, monkeypatch, failure_phase): manager, device, cleanup = graph_case cache = make_cache(device, 1.0) register(manager, cache) entry = next(iter(manager.entries.values())) entry_ref, buffer_ref = weakref.ref(entry), weakref.ref(entry.k[0]) del entry original = manager.model._decode_graph_forward def fail_capture(entry): output = original(entry) if failure_phase == "warmup" or torch.cuda.is_current_stream_capturing(): raise torch.OutOfMemoryError("injected capture OOM") return output monkeypatch.setattr(manager.model, "_decode_graph_forward", fail_capture) assert decode(manager, cache) is None assert entry_ref() is None and buffer_ref() is None assert list(manager.entries.values()) == [None] register(manager, cache) assert decode(manager, cache) is None assert len(cleanup) == 1 assert bool((cache[0]["cond"]["key"] == 1).all()) @pytest.mark.cuda @pytest.mark.gpu def test_replay_uses_active_request_and_layout(graph_case): manager, device, _ = graph_case first, second = make_cache(device, 1.0), make_cache(device, 5.0) first_key = weakref.ref(first[0]["cond"]["key"]) with torch.inference_mode(): register(manager, first) # A second live request aliasing the same key is refused: it keeps its # own cache and decodes eagerly while the owner is in flight. register(manager, second) assert decode(manager, second) is None output = decode(manager, first) expected = torch.ones_like(output).sin() + first[0]["cond"]["key"].mean() + first[0]["cond"]["value"].mean() + 2 torch.testing.assert_close(output, expected) # Once the owner's cache is released, the key is free to adopt. del first assert first_key() is None register(manager, second) output = decode(manager, second) expected = ( torch.ones_like(output).sin() + second[0]["cond"]["key"].mean() + second[0]["cond"]["value"].mean() + 2 ) torch.testing.assert_close(output, expected) original = decode(manager, second).clone() register(manager, second, (1, 1, 4)) other = decode(manager, second, (1, 1, 4)) torch.testing.assert_close(other, torch.ones_like(other).sin() + 5 + 6 + 1) torch.testing.assert_close(decode(manager, second), original) assert all(entry.captures == 1 for entry in manager.entries.values()) @pytest.mark.cpu @pytest.mark.parametrize("inference", [False, True]) def test_prefix_reuse_refreshes_every_layer_without_retaining_requests(graph_case, inference): manager, device, _ = graph_case with torch.inference_mode(inference): cache = make_cache(device, 1.0) + make_cache(device, 2.0) register(manager, cache) entry = next(iter(manager.entries.values())) entry.refresh_prefix(cache) versions = [tensor._version for tensor in entry.k + entry.v] entry.refresh_prefix(cache) assert versions == [tensor._version for tensor in entry.k + entry.v] cache[1]["cond"]["value"] = torch.full_like(cache[1]["cond"]["value"], 7.0) entry.refresh_prefix(cache) assert bool((entry.v[1] == 7).all()) source = weakref.ref(cache[0]["cond"]["key"]) del cache assert source() is None @pytest.mark.cpu def test_prefix_reuse_detects_inplace_changes_and_prefill_registration(graph_case): manager, device, _ = graph_case cache = make_cache(device, 1.0) register(manager, cache) entry = next(iter(manager.entries.values())) entry.refresh_prefix(cache) cache[0]["cond"]["value"].add_(3) entry.refresh_prefix(cache) assert bool((entry.v[0] == 5).all()) register(manager, cache) assert not entry.prefix_sources @pytest.mark.cpu def test_all_valid_mask_shares_graph_entry_and_padding_falls_back(graph_case): manager, device, _ = graph_case cache = make_cache(device, 1.0) register(manager, cache) entry = next(iter(manager.entries.values())) for valid in (True, False): manager.register_prefill( kv_cache=cache, cache_branch="cond", prefix_len=2, img_shapes=[(1, 2, 2)], target_freqs=torch.ones(4, 4), joint_key_valid=torch.full((1, 6), valid), dtype=torch.float32, ) assert list(manager.entries.values()) == [entry] assert entry.attn_metadata is None @pytest.mark.cuda @pytest.mark.gpu def test_batch_changes_and_padding_do_not_reuse_stale_prefix(graph_case): manager, device, _ = graph_case first = make_cache(device, 1.0) second = make_cache(device, 4.0) merged = [ {"cond": {part: torch.cat([first[0]["cond"][part], second[0]["cond"][part]]) for part in ("key", "value")}} ] with torch.inference_mode(): register(manager, first) single = decode(manager, first).clone() register(manager, merged) kwargs = dict( hidden_states=torch.ones(2, 4, 4, device=device), timestep=torch.ones(2, device=device), kv_cache=merged, cache_branch="cond", img_shapes=[(1, 2, 2)], img_mask=torch.tensor([[False, False, True]] * 2, device=device), encoder_hidden_states_mask=torch.ones(2, 3, dtype=torch.bool, device=device), ) result = manager.try_decode(**kwargs) assert result is not None and result.shape == (2, 4, 4) expected = torch.ones_like(result).sin() + torch.tensor([5.0, 11.0], device=device)[:, None, None] torch.testing.assert_close(result, expected) # A rebuilt cache (e.g. reordered rows) is a different live cache # aliasing the same key: decode refuses the shared entry until the # previous owner's tensors are released. flipped = [{"cond": {part: tensor.flip(0) for part, tensor in merged[0]["cond"].items()}}] kwargs["kv_cache"] = flipped assert manager.try_decode(**kwargs) is None del merged register(manager, flipped) torch.testing.assert_close(manager.try_decode(**kwargs), expected.flip(0)) kwargs["encoder_hidden_states_mask"][1, 0] = False assert manager.try_decode(**kwargs) is None torch.testing.assert_close(decode(manager, first), single, rtol=0, atol=0)