diff --git a/tests/diffusion/cache/test_prompt_embed_cache.py b/tests/diffusion/cache/test_prompt_embed_cache.py new file mode 100644 index 000000000..184279c99 --- /dev/null +++ b/tests/diffusion/cache/test_prompt_embed_cache.py @@ -0,0 +1,529 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Unit tests for prompt-embedding cache for diffusion pipelines. + +Covers the module :mod:`vllm_omni.diffusion.cache.prompt_embed_cache`: + +* :class:`PromptEmbedCache` LRU semantics, stats and thread-safety. +* :func:`_hashable` safe-input classification (scalars, torch dtype/device, + numpy scalars, nested containers, tensors/PIL bypass). +* :func:`install_prompt_embed_cache` / :func:`uninstall_prompt_embed_cache` + behaviour: caching, bypass on tensors, bypass on precomputed embeds, + idempotent install, restore on uninstall, detachment of cached tensors. +* :func:`resolve_prompt_embed_cache_config` env-var overrides. +""" + +from __future__ import annotations + +import threading + +import pytest +import torch + +from vllm_omni.diffusion.cache.prompt_embed_cache import ( + _CACHE_MISS, + _NOT_HASHABLE, + PromptEmbedCache, + _build_key, + _detach_output, + _hashable, + install_prompt_embed_cache, + resolve_prompt_embed_cache_config, + uninstall_prompt_embed_cache, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# --------------------------------------------------------------------------- +# _hashable +# --------------------------------------------------------------------------- + + +class TestHashable: + """Tests for the ``_hashable`` helper.""" + + @pytest.mark.parametrize( + "value", + [None, "prompt", "", 0, 1, -5, 3.14, True, False, b"bytes"], + ) + def test_scalars_are_returned_unchanged(self, value): + assert _hashable(value) == value + + def test_torch_device_is_normalized(self): + out = _hashable(torch.device("cpu")) + assert out == ("__torch_device__", "cpu") + assert hash(out) is not None # must actually be hashable + + def test_torch_dtype_is_normalized(self): + out = _hashable(torch.float32) + assert out[0] == "__torch_dtype__" + assert "float32" in out[1] + assert hash(out) is not None + + def test_numpy_scalar_is_unwrapped(self): + np = pytest.importorskip("numpy") + assert _hashable(np.int64(7)) == 7 + assert _hashable(np.float32(1.5)) == pytest.approx(1.5) + + def test_numpy_ndarray_is_not_hashable(self): + np = pytest.importorskip("numpy") + assert _hashable(np.array([1, 2, 3])) is _NOT_HASHABLE + + def test_nested_list_and_tuple(self): + out = _hashable(["a", ("b", 1), [2, 3]]) + assert out[0] == "__seq__" + # fully hashable + assert hash(out) is not None + + def test_dict_is_normalized_by_sorted_keys(self): + a = _hashable({"b": 1, "a": 2}) + b = _hashable({"a": 2, "b": 1}) + assert a == b + assert hash(a) is not None + + def test_tensor_is_not_hashable(self): + assert _hashable(torch.zeros(2)) is _NOT_HASHABLE + + def test_list_containing_tensor_is_not_hashable(self): + assert _hashable(["a", torch.zeros(1)]) is _NOT_HASHABLE + + def test_dict_containing_tensor_is_not_hashable(self): + assert _hashable({"x": torch.zeros(1)}) is _NOT_HASHABLE + + def test_custom_object_is_not_hashable(self): + class Obj: + pass + + assert _hashable(Obj()) is _NOT_HASHABLE + + +# --------------------------------------------------------------------------- +# _detach_output +# --------------------------------------------------------------------------- + + +class TestDetachOutput: + def test_detaches_tensor(self): + t = torch.zeros(2, requires_grad=True) + out = _detach_output(t) + assert isinstance(out, torch.Tensor) + assert not out.requires_grad + + def test_detaches_inside_tuple(self): + t = torch.ones(1, requires_grad=True) + out = _detach_output((t, None)) + assert isinstance(out, tuple) + assert not out[0].requires_grad + assert out[1] is None + + def test_detaches_inside_list_and_dict(self): + t = torch.ones(1, requires_grad=True) + out = _detach_output({"a": [t]}) + assert not out["a"][0].requires_grad + + def test_preserves_namedtuple_type(self): + from collections import namedtuple + + NT = namedtuple("NT", ["x", "y"]) + t = torch.ones(1, requires_grad=True) + out = _detach_output(NT(x=t, y=1)) + assert isinstance(out, NT) + assert not out.x.requires_grad + assert out.y == 1 + + def test_passthrough_for_non_tensor(self): + assert _detach_output("hi") == "hi" + assert _detach_output(None) is None + assert _detach_output(5) == 5 + + +# --------------------------------------------------------------------------- +# PromptEmbedCache +# --------------------------------------------------------------------------- + + +class TestPromptEmbedCache: + def test_rejects_non_positive_max_size(self): + with pytest.raises(ValueError): + PromptEmbedCache(max_size=0) + with pytest.raises(ValueError): + PromptEmbedCache(max_size=-1) + + def test_miss_returns_sentinel_and_bumps_miss_counter(self): + cache = PromptEmbedCache(max_size=2) + assert cache.get("k") is _CACHE_MISS + assert cache.stats()["misses"] == 1 + assert cache.stats()["hits"] == 0 + + def test_put_then_get_returns_value_and_bumps_hit_counter(self): + cache = PromptEmbedCache(max_size=2) + cache.put("k", 42) + assert cache.get("k") == 42 + assert cache.stats()["hits"] == 1 + + def test_none_is_a_cacheable_value(self): + cache = PromptEmbedCache(max_size=2) + cache.put("k", None) + # Distinguishable from a miss. + assert cache.get("k") is None + assert cache.stats()["hits"] == 1 + + def test_lru_eviction_keeps_recently_used(self): + cache = PromptEmbedCache(max_size=2) + cache.put("a", 1) + cache.put("b", 2) + # Access "a" to mark it MRU. + assert cache.get("a") == 1 + cache.put("c", 3) # should evict "b" + assert cache.get("b") is _CACHE_MISS + assert cache.get("a") == 1 + assert cache.get("c") == 3 + assert cache.stats()["size"] == 2 + + def test_get_refreshes_recency(self): + cache = PromptEmbedCache(max_size=2) + cache.put("a", 1) + cache.put("b", 2) + # Hit on "a" → "b" becomes LRU. + cache.get("a") + cache.put("c", 3) + assert cache.get("b") is _CACHE_MISS + assert cache.get("a") == 1 + + def test_put_overwrites_existing_key(self): + cache = PromptEmbedCache(max_size=2) + cache.put("k", 1) + cache.put("k", 2) + assert cache.get("k") == 2 + assert cache.stats()["size"] == 1 + + def test_clear_empties_store_but_keeps_counters(self): + cache = PromptEmbedCache(max_size=2) + cache.put("k", 1) + cache.get("k") + cache.clear() + assert cache.stats()["size"] == 0 + # clear only drops entries, counters reflect prior activity. + assert cache.stats()["hits"] == 1 + + def test_stats_reports_expected_fields(self): + cache = PromptEmbedCache(max_size=4) + stats = cache.stats() + assert set(stats.keys()) == {"size", "max_size", "hits", "misses", "bypassed"} + assert stats["max_size"] == 4 + assert stats["size"] == 0 + + def test_thread_safety_under_concurrent_writes(self): + cache = PromptEmbedCache(max_size=1024) + + def worker(start: int) -> None: + for i in range(start, start + 200): + cache.put(f"k{i}", i) + cache.get(f"k{i}") + + threads = [threading.Thread(target=worker, args=(i * 200,)) for i in range(4)] + for t in threads: + t.start() + for t in threads: + t.join() + + stats = cache.stats() + # 4 * 200 = 800 writes, each followed by a successful read. + assert stats["hits"] == 800 + assert stats["size"] == 800 + + +# --------------------------------------------------------------------------- +# _build_key +# --------------------------------------------------------------------------- + + +def _make_sig(fn): + import inspect + + return inspect.signature(fn) + + +class TestBuildKey: + def test_positional_and_keyword_collide(self): + def encode(prompt, device=None): + return None + + sig = _make_sig(encode) + k1 = _build_key(sig, "m", ("hello",), {"device": "cpu"}) + k2 = _build_key(sig, "m", (), {"prompt": "hello", "device": "cpu"}) + assert k1 == k2 + assert k1 is not None + + def test_different_prompts_produce_different_keys(self): + def encode(prompt): + return None + + sig = _make_sig(encode) + assert _build_key(sig, "m", ("a",), {}) != _build_key(sig, "m", ("b",), {}) + + def test_different_model_tags_produce_different_keys(self): + def encode(prompt): + return None + + sig = _make_sig(encode) + assert _build_key(sig, "m1", ("a",), {}) != _build_key(sig, "m2", ("a",), {}) + + def test_tensor_argument_bypasses(self): + def encode(prompt, extra=None): + return None + + sig = _make_sig(encode) + assert _build_key(sig, "m", ("hi",), {"extra": torch.zeros(1)}) is None + + def test_precomputed_prompt_embeds_bypasses(self): + def encode(prompt=None, prompt_embeds=None): + return None + + sig = _make_sig(encode) + # None → not a bypass. + assert _build_key(sig, "m", (), {"prompt": "hi"}) is not None + # A non-None precomputed embed triggers bypass even if it would + # otherwise be "hashable". + assert _build_key(sig, "m", (), {"prompt": "hi", "prompt_embeds": "x"}) is None + + def test_invalid_binding_returns_none(self): + def encode(prompt): + return None + + sig = _make_sig(encode) + # Unknown kwarg → TypeError on bind → bypass. + assert _build_key(sig, "m", (), {"nope": 1}) is None + + +# --------------------------------------------------------------------------- +# install_prompt_embed_cache / uninstall_prompt_embed_cache +# --------------------------------------------------------------------------- + + +class _FakePipeline: + """Minimal pipeline with a signature-bearing ``encode_prompt``.""" + + def __init__(self): + self.call_count = 0 + + def encode_prompt( + self, + prompt: str, + device=None, + num_images_per_prompt: int = 1, + negative_prompt: str | None = None, + prompt_embeds=None, + negative_prompt_embeds=None, + ): + self.call_count += 1 + # Return a fresh tensor so identity changes between real calls. + return torch.tensor([float(self.call_count)]) + + +class TestInstallAndUninstall: + def test_returns_none_when_pipeline_has_no_encode_prompt(self): + class Empty: + pass + + assert install_prompt_embed_cache(Empty()) is None + + def test_cache_hits_on_identical_args(self): + pipe = _FakePipeline() + cache = install_prompt_embed_cache(pipe, max_size=4) + assert cache is not None + + out1 = pipe.encode_prompt("cat", device="cpu") + out2 = pipe.encode_prompt("cat", device="cpu") + # Second call must be served from cache (wrapped fn not invoked). + assert pipe.call_count == 1 + # Cache returns the same tensor object. + assert torch.equal(out1, out2) + stats = cache.stats() + assert stats["hits"] == 1 + assert stats["misses"] == 1 + assert stats["bypassed"] == 0 + + def test_positional_and_keyword_share_cache_slot(self): + pipe = _FakePipeline() + cache = install_prompt_embed_cache(pipe) + pipe.encode_prompt("cat") + pipe.encode_prompt(prompt="cat") + assert pipe.call_count == 1 + assert cache.stats()["hits"] == 1 + + def test_different_prompts_miss(self): + pipe = _FakePipeline() + cache = install_prompt_embed_cache(pipe) + pipe.encode_prompt("cat") + pipe.encode_prompt("dog") + assert pipe.call_count == 2 + assert cache.stats()["misses"] == 2 + + def test_tensor_argument_bypasses_and_increments_counter(self): + pipe = _FakePipeline() + cache = install_prompt_embed_cache(pipe) + # Put a tensor into a non-precomputed-embed slot to trigger bypass. + pipe.encode_prompt("cat", device=torch.device("cpu"), num_images_per_prompt=1) + + # Now pass a fake tensor via a harmless positional-like call: use + # the negative_prompt slot with an unhashable object. + class Unhashable: + pass + + pipe.encode_prompt("cat", negative_prompt=Unhashable()) # type: ignore[arg-type] + pipe.encode_prompt("cat", negative_prompt=Unhashable()) # type: ignore[arg-type] + assert cache.stats()["bypassed"] == 2 + # All three calls actually executed the underlying function. + assert pipe.call_count == 3 + + def test_precomputed_embeds_bypass_cache(self): + pipe = _FakePipeline() + cache = install_prompt_embed_cache(pipe) + pipe.encode_prompt("cat", prompt_embeds=torch.zeros(1)) + pipe.encode_prompt("cat", prompt_embeds=torch.zeros(1)) + assert cache.stats()["bypassed"] == 2 + assert cache.stats()["hits"] == 0 + assert pipe.call_count == 2 + + def test_install_is_idempotent(self): + pipe = _FakePipeline() + c1 = install_prompt_embed_cache(pipe, max_size=4) + c2 = install_prompt_embed_cache(pipe, max_size=16) + assert c1 is c2 + # The second call must not re-wrap (max_size stays from first install). + assert c1.max_size == 4 + + def test_disabled_cache_passes_through(self): + pipe = _FakePipeline() + cache = install_prompt_embed_cache(pipe, enabled=False) + pipe.encode_prompt("cat") + pipe.encode_prompt("cat") + assert pipe.call_count == 2 + assert cache.stats()["hits"] == 0 + assert cache.stats()["misses"] == 0 + + def test_toggle_enabled_flag_takes_effect(self): + pipe = _FakePipeline() + cache = install_prompt_embed_cache(pipe, enabled=True) + pipe.encode_prompt("cat") + pipe.encode_prompt("cat") + assert cache.stats()["hits"] == 1 + + cache.enabled = False + pipe.encode_prompt("cat") + pipe.encode_prompt("cat") + # Cache-bypassed calls go through; hit count unchanged. + assert cache.stats()["hits"] == 1 + + def test_cached_tensor_is_detached(self): + class GradPipeline: + def encode_prompt(self, prompt): + return torch.ones(1, requires_grad=True) + + pipe = GradPipeline() + install_prompt_embed_cache(pipe) + # First call populates the cache. + pipe.encode_prompt("cat") + # Second call returns the cached tensor which must be detached. + cached = pipe.encode_prompt("cat") + assert isinstance(cached, torch.Tensor) + assert not cached.requires_grad + + def test_lru_eviction_through_wrapper(self): + pipe = _FakePipeline() + cache = install_prompt_embed_cache(pipe, max_size=2) + pipe.encode_prompt("a") + pipe.encode_prompt("b") + pipe.encode_prompt("c") # evicts "a" + pipe.encode_prompt("a") # miss again → re-run + assert pipe.call_count == 4 + assert cache.stats()["size"] == 2 + + def test_uninstall_restores_original(self): + pipe = _FakePipeline() + original = pipe.encode_prompt + install_prompt_embed_cache(pipe) + assert pipe.encode_prompt is not original + uninstall_prompt_embed_cache(pipe) + # The bound method reference will differ, but identity of underlying + # function matches and wrapper attributes are gone. + assert not hasattr(pipe, "_prompt_embed_cache") or pipe._prompt_embed_cache is None + # After uninstall calls go straight to the real function. + pipe.encode_prompt("cat") + pipe.encode_prompt("cat") + assert pipe.call_count == 2 + + def test_uninstall_on_unwrapped_pipeline_is_noop(self): + class Empty: + def encode_prompt(self, prompt): + return None + + pipe = Empty() + uninstall_prompt_embed_cache(pipe) # must not raise + + +# --------------------------------------------------------------------------- +# resolve_prompt_embed_cache_config +# --------------------------------------------------------------------------- + + +@pytest.fixture +def clean_env(monkeypatch): + monkeypatch.delenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raising=False) + monkeypatch.delenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", raising=False) + + +class TestResolveConfig: + def test_defaults_when_nothing_set(self, clean_env): + enable, size = resolve_prompt_embed_cache_config() + assert enable is False + assert size == 32 + + def test_explicit_args_used(self, clean_env): + enable, size = resolve_prompt_embed_cache_config(enable=True, max_size=8) + assert enable is True + assert size == 8 + + @pytest.mark.parametrize("raw", ["1", "true", "TRUE", "yes", "on"]) + def test_env_enable_truthy(self, clean_env, monkeypatch, raw): + monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raw) + enable, _ = resolve_prompt_embed_cache_config(enable=False) + assert enable is True + + @pytest.mark.parametrize("raw", ["0", "false", "FALSE", "no", "off"]) + def test_env_enable_falsy_overrides_true(self, clean_env, monkeypatch, raw): + monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raw) + enable, _ = resolve_prompt_embed_cache_config(enable=True) + assert enable is False + + def test_env_enable_invalid_leaves_arg_untouched(self, clean_env, monkeypatch): + monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", "garbage") + enable, _ = resolve_prompt_embed_cache_config(enable=True) + assert enable is True + + def test_env_size_overrides_arg(self, clean_env, monkeypatch): + monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "7") + _, size = resolve_prompt_embed_cache_config(max_size=99) + assert size == 7 + + def test_env_size_zero_or_negative_ignored(self, clean_env, monkeypatch): + monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "0") + _, size = resolve_prompt_embed_cache_config(max_size=12) + assert size == 12 + + monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "-5") + _, size = resolve_prompt_embed_cache_config(max_size=12) + assert size == 12 + + def test_env_size_non_integer_ignored(self, clean_env, monkeypatch): + monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "notanint") + _, size = resolve_prompt_embed_cache_config(max_size=12) + assert size == 12 + + def test_max_size_none_falls_back_to_default(self, clean_env): + _, size = resolve_prompt_embed_cache_config(enable=True, max_size=None) + assert size == 32 diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index cb12019bb..4205d8280 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -52,6 +52,7 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar runner.pipeline = _DummyPipeline(output=SimpleNamespace(output="ok")) runner.cache_backend = cache_backend runner.offload_backend = None + runner.prompt_embed_cache = None runner.od_config = SimpleNamespace( cache_backend=cache_backend_name, enable_cache_dit_summary=enable_cache_dit_summary, diff --git a/vllm_omni/diffusion/cache/__init__.py b/vllm_omni/diffusion/cache/__init__.py index a5968f612..1d76a8262 100644 --- a/vllm_omni/diffusion/cache/__init__.py +++ b/vllm_omni/diffusion/cache/__init__.py @@ -11,6 +11,12 @@ Cache backends are instantiated directly via their constructors and configured v """ from vllm_omni.diffusion.cache.base import CacheBackend +from vllm_omni.diffusion.cache.prompt_embed_cache import ( + PromptEmbedCache, + install_prompt_embed_cache, + resolve_prompt_embed_cache_config, + uninstall_prompt_embed_cache, +) from vllm_omni.diffusion.cache.teacache import ( CacheContext, TeaCacheConfig, @@ -24,4 +30,8 @@ __all__ = [ "CacheContext", "TeaCacheBackend", "apply_teacache_hook", + "PromptEmbedCache", + "install_prompt_embed_cache", + "resolve_prompt_embed_cache_config", + "uninstall_prompt_embed_cache", ] diff --git a/vllm_omni/diffusion/cache/prompt_embed_cache.py b/vllm_omni/diffusion/cache/prompt_embed_cache.py new file mode 100644 index 000000000..f95197a30 --- /dev/null +++ b/vllm_omni/diffusion/cache/prompt_embed_cache.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Prompt-embedding cache for diffusion pipelines. + +Text encoders are the single most expensive per-request preprocessing step in +most diffusion pipelines, and their inputs are frequently identical across +requests (e.g. in GRPO-style rollouts the same prompt is submitted many times +with different seeds to sample different images). Re-running the text encoder +for each of those requests wastes compute and GPU memory. + +This module provides a small LRU cache plus a transparent wrapper for a +pipeline's ``encode_prompt`` method. Because almost every diffusion pipeline in +``vllm_omni/diffusion/models`` routes all text-encoder invocations through +``self.encode_prompt(...)``, wrapping that single method is sufficient to +cache results model-wide without per-pipeline edits. + +Design points: + * The wrapper is installed by :class:`DiffusionModelRunner` after the + pipeline has loaded, so each runner process owns its own cache. + * Cache keys are derived from the bound ``encode_prompt`` arguments. Only + inputs we can safely hash (``str`` / ``int`` / ``float`` / ``bool`` / + ``None`` / ``bytes`` / ``torch.device`` / ``torch.dtype`` / numpy + scalars / nested lists/tuples/dicts of those) participate in the key. + If any argument is a tensor, PIL image, or other non-trivial object, we + bypass the cache for that call to guarantee correctness. + * If a caller passes precomputed ``*_embeds`` into ``encode_prompt`` the + wrapper also bypasses the cache because the call is already short-circuit. + * Cache values are detached tensors. Downstream pipeline code typically + does non-inplace ops (``.repeat``, ``.view``, slicing); we do not clone + on hit since outputs are treated as read-only. +""" + +from __future__ import annotations + +import functools +import inspect +import os +import threading +from collections import OrderedDict +from typing import Any + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# Sentinel returned by :func:`_hashable` when an input cannot be hashed +# safely (e.g. tensor / PIL image / custom object). Any such input disables +# caching for that specific call. +_NOT_HASHABLE = object() + +# Sentinel returned by :meth:`PromptEmbedCache.get` to distinguish a cache +# miss from a legitimate cached value of ``None``. +_CACHE_MISS = object() + +# Argument names whose presence of a non-None value means the caller is +# providing precomputed embeddings. In every pipeline we inspected, passing +# these makes ``encode_prompt`` a cheap passthrough, so caching it has no +# benefit and just wastes memory on large tensor keys. +_PRECOMPUTED_EMBED_ARGS = ( + "prompt_embeds", + "negative_prompt_embeds", + "pooled_prompt_embeds", + "negative_pooled_prompt_embeds", + "prompt_embeds_mask", + "negative_prompt_embeds_mask", +) + + +def _hashable(obj: Any) -> Any: + """Convert *obj* into a hashable representation or return ``_NOT_HASHABLE``. + + Only inputs whose textual / scalar identity fully determines the + ``encode_prompt`` output are considered safe. Tensors and PIL images (which + flow through e.g. image-edit pipelines) deliberately disable caching for + that call rather than risk stale results. Common value-like configuration + objects such as ``torch.device`` and ``torch.dtype`` are normalized to + stable string forms so they can participate in cache keys safely. + """ + if obj is None or isinstance(obj, (str, int, float, bool, bytes)): + return obj + if isinstance(obj, torch.device): + return ("__torch_device__", str(obj)) + if isinstance(obj, torch.dtype): + return ("__torch_dtype__", str(obj)) + if type(obj).__module__.split(".", 1)[0] == "numpy" and hasattr(obj, "item"): + try: + return _hashable(obj.item()) + except (ValueError, TypeError): + return _NOT_HASHABLE + if isinstance(obj, (list, tuple)): + out = [] + for item in obj: + h = _hashable(item) + if h is _NOT_HASHABLE: + return _NOT_HASHABLE + out.append(h) + return ("__seq__", tuple(out)) + if isinstance(obj, dict): + items = [] + for k in sorted(obj.keys(), key=repr): + v = _hashable(obj[k]) + if v is _NOT_HASHABLE: + return _NOT_HASHABLE + items.append((repr(k), v)) + return ("__dict__", tuple(items)) + # torch.Tensor, PIL.Image, np.ndarray, and any other object → not safe. + return _NOT_HASHABLE + + +def _detach_output(value: Any) -> Any: + """Detach tensors in the return value so the cache never holds grad state.""" + if isinstance(value, torch.Tensor): + return value.detach() + if isinstance(value, tuple): + detached = tuple(_detach_output(v) for v in value) + # Preserve namedtuple types so attribute access keeps working. + if hasattr(value, "_fields"): + return type(value)(*detached) + return detached + if isinstance(value, list): + return [_detach_output(v) for v in value] + if isinstance(value, dict): + return {k: _detach_output(v) for k, v in value.items()} + return value + + +class PromptEmbedCache: + """Thread-safe LRU cache for ``encode_prompt`` outputs. + + The cache stores whatever ``encode_prompt`` returns (tensor, tuple of + tensors, ``None``, etc.). Lookup / insertion is O(1) amortized. Eviction + is least-recently-used. + """ + + def __init__(self, max_size: int = 32, enabled: bool = True) -> None: + if max_size <= 0: + raise ValueError(f"max_size must be positive, got {max_size}") + self.max_size = max_size + self.enabled = enabled + self._store: OrderedDict[Any, Any] = OrderedDict() + self._lock = threading.Lock() + self.hits = 0 + self.misses = 0 + self.bypassed = 0 + + def get(self, key: Any) -> Any: + """Return the cached value or ``_CACHE_MISS`` if absent. + + Returning a sentinel (rather than ``None``) lets callers cache + legitimate ``None`` results from the wrapped function. + """ + with self._lock: + if key not in self._store: + self.misses += 1 + return _CACHE_MISS + self._store.move_to_end(key) + self.hits += 1 + return self._store[key] + + def put(self, key: Any, value: Any) -> None: + with self._lock: + self._store[key] = value + self._store.move_to_end(key) + while len(self._store) > self.max_size: + self._store.popitem(last=False) + + def clear(self) -> None: + with self._lock: + self._store.clear() + + def stats(self) -> dict[str, int]: + with self._lock: + return { + "size": len(self._store), + "max_size": self.max_size, + "hits": self.hits, + "misses": self.misses, + "bypassed": self.bypassed, + } + + +def _build_key( + signature: inspect.Signature, + model_tag: str, + args: tuple, + kwargs: dict, +) -> Any: + """Return a hashable cache key or ``None`` if this call must bypass cache. + + We bind the raw positional / keyword arguments to the method signature so + that keyword and positional invocations with the same underlying values + collide in the cache. + """ + try: + bound = signature.bind_partial(*args, **kwargs) + except TypeError: + return None + bound.apply_defaults() + + # Bypass when the caller supplies precomputed embeddings. + for name in _PRECOMPUTED_EMBED_ARGS: + if name in bound.arguments and bound.arguments[name] is not None: + return None + + items: list[tuple[str, Any]] = [] + for name in sorted(bound.arguments.keys()): + h = _hashable(bound.arguments[name]) + if h is _NOT_HASHABLE: + return None + items.append((name, h)) + return (model_tag, tuple(items)) + + +def install_prompt_embed_cache( + pipeline: Any, + *, + max_size: int = 32, + enabled: bool = True, + model_tag: str | None = None, +) -> PromptEmbedCache | None: + """Wrap ``pipeline.encode_prompt`` so results are cached by argument identity. + + Idempotent: calling twice on the same pipeline is a no-op and returns the + existing cache. Returns ``None`` if the pipeline has no ``encode_prompt`` + method. + """ + existing = getattr(pipeline, "_prompt_embed_cache", None) + if existing is not None: + return existing + + encode_fn = getattr(pipeline, "encode_prompt", None) + if encode_fn is None or not callable(encode_fn): + logger.debug( + "Prompt-embed cache: pipeline %s has no encode_prompt(); skipping.", + type(pipeline).__name__, + ) + return None + + try: + signature = inspect.signature(encode_fn) + except (TypeError, ValueError) as e: + logger.warning( + "Prompt-embed cache: cannot introspect encode_prompt on %s (%s); skipping.", + type(pipeline).__name__, + e, + ) + return None + + cache = PromptEmbedCache(max_size=max_size, enabled=enabled) + tag = model_tag or type(pipeline).__name__ + + @functools.wraps(encode_fn) + def cached_encode_prompt(*args: Any, **kwargs: Any) -> Any: + if not cache.enabled: + return encode_fn(*args, **kwargs) + key = _build_key(signature, tag, args, kwargs) + if key is None: + with cache._lock: + cache.bypassed += 1 + return encode_fn(*args, **kwargs) + hit = cache.get(key) + if hit is not _CACHE_MISS: + return hit + out = encode_fn(*args, **kwargs) + cache.put(key, _detach_output(out)) + return out + + cached_encode_prompt.__wrapped__ = encode_fn # type: ignore[attr-defined] + # Use object.__setattr__ to bypass diffusers ModelMixin attribute checks. + try: + pipeline.encode_prompt = cached_encode_prompt # type: ignore[assignment] + except Exception: + object.__setattr__(pipeline, "encode_prompt", cached_encode_prompt) + object.__setattr__(pipeline, "_prompt_embed_cache", cache) + object.__setattr__(pipeline, "_prompt_embed_cache_original_fn", encode_fn) + + logger.info( + "Prompt-embed cache installed on %s (max_size=%d, enabled=%s).", + tag, + max_size, + enabled, + ) + return cache + + +def uninstall_prompt_embed_cache(pipeline: Any) -> None: + """Restore the original ``encode_prompt`` on *pipeline* if wrapped.""" + original = getattr(pipeline, "_prompt_embed_cache_original_fn", None) + if original is None: + return + try: + pipeline.encode_prompt = original # type: ignore[assignment] + except Exception: + object.__setattr__(pipeline, "encode_prompt", original) + for attr in ("_prompt_embed_cache", "_prompt_embed_cache_original_fn"): + if hasattr(pipeline, attr): + try: + delattr(pipeline, attr) + except Exception: + object.__setattr__(pipeline, attr, None) + + +def resolve_prompt_embed_cache_config( + enable: bool | None = None, + max_size: int | None = None, +) -> tuple[bool, int]: + """Combine explicit args with env-var overrides. + + Environment variables (useful for quick enablement in GRPO jobs without + touching config files): + + ``OMNI_DIFFUSION_PROMPT_EMBED_CACHE`` (``1``/``0``/``true``/``false``) + ``OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE`` (positive int) + """ + env_enable = os.environ.get("OMNI_DIFFUSION_PROMPT_EMBED_CACHE") + if env_enable is not None: + parsed = env_enable.strip().lower() + if parsed in ("1", "true", "yes", "on"): + enable = True + elif parsed in ("0", "false", "no", "off"): + enable = False + + env_size = os.environ.get("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE") + if env_size is not None: + try: + env_size_int = int(env_size) + if env_size_int > 0: + max_size = env_size_int + except ValueError: + logger.warning( + "Ignoring non-integer OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE=%r.", + env_size, + ) + + return bool(enable), int(max_size or 32) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 433ea5577..b0b3287a3 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -432,6 +432,15 @@ class OmniDiffusionConfig: cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict) enable_cache_dit_summary: bool = False + # Prompt-embedding cache. When enabled, ``DiffusionModelRunner`` wraps the + # pipeline's ``encode_prompt`` so repeated calls with identical prompt + # arguments (e.g. GRPO rollouts that sample the same prompt many times + # with different seeds) reuse the text-encoder output instead of re-running + # it. Safe against inputs that cannot be hashed (tensors, PIL images): + # those calls transparently bypass the cache. + enable_prompt_embed_cache: bool = False + prompt_embed_cache_size: int = 32 + # Distributed executor backend distributed_executor_backend: str = "mp" nccl_port: int | None = None diff --git a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py index e223c516e..47b905de7 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py +++ b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py @@ -446,6 +446,33 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor )[0].squeeze(0) return video_noise, audio_noise + def encode_prompt( + self, + prompt: str, + video_negative_prompt: str = "", + audio_negative_prompt: str = "", + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Encode the positive and negative prompts via ``self.text_model``. + + This is the single text-encoder entrypoint so that the runner-level + prompt-embedding cache (see + ``vllm_omni/diffusion/cache/prompt_embed_cache.py``) can transparently + memoize results when the same prompts are submitted repeatedly (e.g. + GRPO rollouts that sample the same prompt with different seeds). + + Returns: + (audio_pos, video_pos, video_neg, audio_neg) embeddings cast to + ``self.target_dtype``. ``audio_pos`` and ``video_pos`` are the + same tensor, matching the original inline behavior. + """ + text_embeddings = self.text_model( + [prompt, video_negative_prompt, audio_negative_prompt], + device=self.device, + ) + text_embeddings = [emb.to(self.target_dtype) for emb in text_embeddings] + # Index 0 is the positive embedding, shared between video and audio branches. + return text_embeddings[0], text_embeddings[0], text_embeddings[1], text_embeddings[2] + def forward( self, request: OmniDiffusionRequest, @@ -503,12 +530,13 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor ) # 3. text embedding - text_embeddings = self.text_model([prompt, video_negative_prompt, audio_negative_prompt], device=self.device) - text_embeddings = [emb.to(self.target_dtype) for emb in text_embeddings] - text_embeddings_audio_pos = text_embeddings[0] - text_embeddings_video_pos = text_embeddings[0] - text_embeddings_video_neg = text_embeddings[1] - text_embeddings_audio_neg = text_embeddings[2] + text_embeddings_audio_pos, text_embeddings_video_pos, text_embeddings_video_neg, text_embeddings_audio_neg = ( + self.encode_prompt( + prompt=prompt, + video_negative_prompt=video_negative_prompt, + audio_negative_prompt=audio_negative_prompt, + ) + ) video_latent_h, video_latent_w = height // 16, width // 16 diff --git a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py index 7f073ab87..65da889f1 100644 --- a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py +++ b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py @@ -2130,6 +2130,28 @@ class MagiHumanPipeline(nn.Module, ProgressBarMixin, DiffusionPipelineProfilerMi return resample(audio_np, target_len) + def encode_prompt( + self, + prompt: str, + target_length: int | None = None, + ) -> tuple[torch.Tensor, int]: + """Encode *prompt* with the T5-Gemma text encoder and pad to fixed length. + + This is the single text-encoder entrypoint so the runner-level + prompt-embedding cache (see + ``vllm_omni/diffusion/cache/prompt_embed_cache.py``) can transparently + memoize results when the same prompt is submitted repeatedly. + + Returns: + ``(context, original_context_len)`` matching + :func:`_get_padded_t5_gemma_embedding`. + """ + return _get_padded_t5_gemma_embedding( + prompt, + self.text_encoder, + target_length if target_length is not None else self.t5_gemma_target_length, + ) + @torch.inference_mode() def forward( self, @@ -2200,11 +2222,7 @@ class MagiHumanPipeline(nn.Module, ProgressBarMixin, DiffusionPipelineProfilerMi device=device, ) - context, original_context_len = _get_padded_t5_gemma_embedding( - prompt, - self.text_encoder, - self.t5_gemma_target_length, - ) + context, original_context_len = self.encode_prompt(prompt) if image_path is not None: br_image = self._encode_image(load_image(image_path), br_height, br_width) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 2720de038..b18d8fe97 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -23,6 +23,10 @@ from vllm.logger import init_logger from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes from vllm_omni.diffusion.cache.cache_dit_backend import cache_summary +from vllm_omni.diffusion.cache.prompt_embed_cache import ( + install_prompt_embed_cache, + resolve_prompt_embed_cache_config, +) from vllm_omni.diffusion.cache.selector import get_cache_backend from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -71,6 +75,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): self.pipeline = None self.cache_backend = None self.offload_backend = None + self.prompt_embed_cache = None # Cache for per-request stepwise state. self.state_cache: dict[str, DiffusionRequestState] = {} @@ -187,12 +192,38 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): else: self.cache_backend.enable(self.pipeline) + # Install prompt-embedding cache (transparent wrapper around + # ``pipeline.encode_prompt``). Enabled via config or env var; a no-op + # when the pipeline does not expose ``encode_prompt``. + enable_pec, pec_size = resolve_prompt_embed_cache_config( + enable=getattr(self.od_config, "enable_prompt_embed_cache", False), + max_size=getattr(self.od_config, "prompt_embed_cache_size", 32), + ) + if enable_pec: + self.prompt_embed_cache = install_prompt_embed_cache( + self.pipeline, + max_size=pec_size, + enabled=True, + model_tag=self.od_config.model_class_name, + ) + logger.info("Model runner: Initialization complete.") def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights into the pipeline.""" return self.pipeline.load_weights(weights) + def clear_prompt_embed_cache(self) -> None: + """Evict all cached text-encoder outputs (e.g. between training epochs).""" + if self.prompt_embed_cache is not None: + self.prompt_embed_cache.clear() + + def get_prompt_embed_cache_stats(self) -> dict | None: + """Return hit/miss statistics for the prompt-embedding cache, if enabled.""" + if self.prompt_embed_cache is None: + return None + return self.prompt_embed_cache.stats() + def _record_peak_memory(self, output: DiffusionOutput) -> None: """Record peak GPU memory for the current forward pass into output. @@ -299,6 +330,10 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): if is_primary: self._record_peak_memory(output) + # Log prompt-embed cache activity (hits/misses accumulate across requests). + if is_primary and self.prompt_embed_cache is not None: + logger.debug("prompt-embed cache: %s", self.prompt_embed_cache.stats()) + # NOTE: if ( self.cache_backend is not None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 4986eae63..440b490b2 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1949,6 +1949,8 @@ class AsyncOmniEngine: "mp" if kwargs.get("distributed_executor_backend") is None else kwargs["distributed_executor_backend"] ), "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), + "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False), + "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32), "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), "quantization": kwargs.get("quantization", None),