[feat] Support prompt embedding caching for diffusion model (#2962)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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

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,

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",
 ]

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)

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

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_audio_pos, text_embeddings_video_pos, text_embeddings_video_neg, text_embeddings_audio_neg = (
-        text_embeddings = [emb.to(self.target_dtype) for emb in text_embeddings]
+            self.encode_prompt(
-        text_embeddings_audio_pos = text_embeddings[0]
+                prompt=prompt,
-        text_embeddings_video_pos = text_embeddings[0]
+                video_negative_prompt=video_negative_prompt,
-        text_embeddings_video_neg = text_embeddings[1]
+                audio_negative_prompt=audio_negative_prompt,
-        text_embeddings_audio_neg = text_embeddings[2]
+            )
+        )
 
         video_latent_h, video_latent_w = height // 16, width // 16
 

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(
+        context, original_context_len = self.encode_prompt(prompt)
-            prompt,
-            self.text_encoder,
-            self.t5_gemma_target_length,
-        )
 
         if image_path is not None:
             br_image = self._encode_image(load_image(image_path), br_height, br_width)

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

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