[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>

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)