246 lines
10 KiB
Python
246 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Midway prompt updates for chunked diffusion step execution.
|
|
|
|
``PromptUpdateMixin`` exposes
|
|
- a pipeline interface for the external to queue a prompt update,
|
|
- a pipeline internal method to apply prompt transitions at chunk boundaries.
|
|
|
|
``prompt_update_versions``: a helper for ``InputBatch`` cache invalidation.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
from threading import Lock
|
|
from typing import ClassVar, TypedDict, cast
|
|
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.worker.utils import StepRequestState
|
|
|
|
DEFAULT_TRANSITION_CHUNKS = 3
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PromptUpdateExtra(TypedDict, total=False):
|
|
"""A "protocol" for ``StepRequestState.extra`` keys used by the prompt-update feature."""
|
|
|
|
pending_prompt_update: _PendingPromptUpdate
|
|
prompt_update_state: _PromptUpdateState
|
|
prompt_update_chunk_metadata: _PromptUpdateChunkMetadata
|
|
prompt_update_version: int # Bumped when prompt_embeds change; invalidates InputBatch cache.
|
|
|
|
# The lock for storing/updating pending prompt updates.
|
|
# Saved in state.extra to be naturally cleaned up with the request state.
|
|
# As PromptUpdateMixin (the pipeline) currently lacks an obvious request-finish hook
|
|
prompt_update_lock: Lock
|
|
|
|
|
|
def prompt_update_versions(states: Sequence[StepRequestState]) -> tuple[int, ...]:
|
|
"""Return per-request prompt-update versions for batch cache comparison.
|
|
|
|
To be used by ``InputBatch`` or other external to determine if the cached batch is still valid.
|
|
"""
|
|
return tuple(int(cast(PromptUpdateExtra, state.extra).get("prompt_update_version", 0)) for state in states)
|
|
|
|
|
|
class PromptUpdateMixin:
|
|
"""Mixin for chunked streaming pipelines that support midway prompt updates.
|
|
|
|
Implements the behavior expected by :class:`~vllm_omni.diffusion.models.interface.SupportsPromptUpdate`.
|
|
All per-request transition state lives on ``StepRequestState.extra``; this mixin itself is stateless.
|
|
"""
|
|
|
|
supports_prompt_update: ClassVar[bool] = True
|
|
|
|
@staticmethod
|
|
def _prompt_update_lock(extra: PromptUpdateExtra) -> Lock:
|
|
return extra.setdefault("prompt_update_lock", Lock())
|
|
|
|
def prepare_prompt_update(
|
|
self,
|
|
state: StepRequestState,
|
|
prompt: str,
|
|
event_id: str,
|
|
transition_chunks: int | None = None,
|
|
) -> None:
|
|
"""Encode and queue a prompt update to apply at the next chunk boundary.
|
|
It does not actually apply the prompt transition.
|
|
|
|
This is an extended interface for the pipeline, to be called by the runner.
|
|
"""
|
|
if not prompt:
|
|
raise ValueError("prompt must be non-empty")
|
|
if not event_id:
|
|
raise ValueError("event_id must be non-empty")
|
|
if state.prompt_embeds is None:
|
|
raise ValueError(
|
|
f"prompt_update is not allowed before initial generation has started (request_id={state.request_id!r})"
|
|
)
|
|
duration = DEFAULT_TRANSITION_CHUNKS if transition_chunks is None else transition_chunks
|
|
if duration < 0:
|
|
raise ValueError("transition_chunks must be >= 0")
|
|
|
|
target_prompt_embeds, _ = self.encode_prompt(
|
|
prompt=prompt,
|
|
negative_prompt=None,
|
|
do_classifier_free_guidance=False,
|
|
num_videos_per_prompt=state.sampling.num_outputs_per_prompt,
|
|
max_sequence_length=state.sampling.max_sequence_length,
|
|
device=self.device,
|
|
dtype=self.transformer.dtype,
|
|
)
|
|
extra = cast(PromptUpdateExtra, state.extra)
|
|
with self._prompt_update_lock(extra):
|
|
extra["pending_prompt_update"] = {
|
|
"event_id": event_id,
|
|
"prompt": prompt,
|
|
"target_prompt_embeds": target_prompt_embeds,
|
|
"transition_chunks": duration,
|
|
}
|
|
|
|
def _apply_prompt_update_at_chunk_boundary(self, state: StepRequestState) -> None:
|
|
"""Advance or start prompt interpolation before the next chunk."""
|
|
extra = cast(PromptUpdateExtra, state.extra)
|
|
update_state: _PromptUpdateState | None = extra.get("prompt_update_state")
|
|
with self._prompt_update_lock(extra):
|
|
pending = extra.pop("pending_prompt_update", None)
|
|
embeds_changed = False
|
|
next_chunk_index = state.chunk_index
|
|
started_event_ids: list[str] = []
|
|
active_event_ids: list[str] = []
|
|
completed_event_ids: list[str] = []
|
|
|
|
# If current transition is not complete, advance it.
|
|
# After completion, leave prompt_update_state in place (so a later pending
|
|
# update can abort/overwrite), but do not keep bumping the version.
|
|
if update_state is not None:
|
|
in_transition = (
|
|
update_state.transition_chunks > 0
|
|
and update_state.elapsed_transition_chunks < update_state.transition_chunks
|
|
)
|
|
if in_transition:
|
|
update_state.advance_transition()
|
|
state.prompt_embeds = update_state.blended_prompt_embeds()
|
|
embeds_changed = True
|
|
active_event_ids.append(update_state.event_id)
|
|
if update_state.elapsed_transition_chunks >= update_state.transition_chunks:
|
|
state.prompt_embeds = update_state.target_prompt_embeds
|
|
update_state.source_prompt_embeds = update_state.target_prompt_embeds
|
|
completed_event_ids.append(update_state.event_id)
|
|
if update_state.target_prompt is not None:
|
|
logger.debug(
|
|
"prompt_update transition complete request_id=%s next_chunk_index=%d prompt=%.20s...",
|
|
state.request_id,
|
|
next_chunk_index,
|
|
update_state.target_prompt,
|
|
)
|
|
|
|
# If a new prompt update is pending, start a new transition.
|
|
if pending is not None:
|
|
if state.prompt_embeds is None:
|
|
raise RuntimeError(
|
|
"internal error: trying to apply a pending prompt update but "
|
|
f"current prompt_embeds is None (request_id={state.request_id!r})"
|
|
)
|
|
source = state.prompt_embeds.detach().clone()
|
|
target = pending["target_prompt_embeds"]
|
|
duration = int(pending["transition_chunks"])
|
|
prompt = str(pending.get("prompt", ""))
|
|
event_id = pending.get("event_id")
|
|
update_state = _PromptUpdateState(
|
|
source_prompt_embeds=source,
|
|
target_prompt_embeds=target,
|
|
transition_chunks=duration,
|
|
event_id=event_id,
|
|
target_prompt=prompt,
|
|
)
|
|
extra["prompt_update_state"] = update_state
|
|
started_event_ids.append(event_id)
|
|
active_event_ids.append(event_id)
|
|
if duration <= 0:
|
|
# A hard/immediate transition
|
|
state.prompt_embeds = target
|
|
completed_event_ids.append(event_id)
|
|
logger.debug(
|
|
"prompt_update sharp transition request_id=%s next_chunk_index=%d prompt=%.20s...",
|
|
state.request_id,
|
|
next_chunk_index,
|
|
prompt,
|
|
)
|
|
else:
|
|
# A transition that really takes some time
|
|
update_state.advance_transition()
|
|
state.prompt_embeds = update_state.blended_prompt_embeds()
|
|
if update_state.elapsed_transition_chunks >= update_state.transition_chunks:
|
|
state.prompt_embeds = update_state.target_prompt_embeds
|
|
update_state.source_prompt_embeds = update_state.target_prompt_embeds
|
|
completed_event_ids.append(event_id)
|
|
logger.debug(
|
|
"prompt_update transition start request_id=%s next_chunk_index=%d prompt=%.20s...",
|
|
state.request_id,
|
|
next_chunk_index,
|
|
prompt,
|
|
)
|
|
embeds_changed = True
|
|
|
|
# Indicate that prompt embeddings have changed---clear input batch cache.
|
|
if embeds_changed:
|
|
new_version = int(extra.get("prompt_update_version", 0)) + 1
|
|
extra["prompt_update_version"] = new_version
|
|
|
|
extra["prompt_update_chunk_metadata"] = {
|
|
"started_event_ids": started_event_ids,
|
|
"active_event_ids": active_event_ids,
|
|
"completed_event_ids": completed_event_ids,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class _PromptUpdateState:
|
|
"""Per-request prompt interpolation stored in ``state.extra["prompt_update_state"]``."""
|
|
|
|
source_prompt_embeds: torch.Tensor
|
|
target_prompt_embeds: torch.Tensor
|
|
transition_chunks: int
|
|
event_id: str
|
|
elapsed_transition_chunks: int = 0
|
|
target_prompt: str | None = None
|
|
|
|
def blended_prompt_embeds(self) -> torch.Tensor:
|
|
alpha = self._current_alpha()
|
|
if alpha >= 1.0:
|
|
return self.target_prompt_embeds
|
|
return (1.0 - alpha) * self.source_prompt_embeds + alpha * self.target_prompt_embeds
|
|
|
|
def advance_transition(self) -> None:
|
|
if self.transition_chunks <= 0:
|
|
self.source_prompt_embeds = self.target_prompt_embeds
|
|
return
|
|
self.elapsed_transition_chunks += 1
|
|
if self.elapsed_transition_chunks >= self.transition_chunks:
|
|
self.source_prompt_embeds = self.target_prompt_embeds
|
|
self.elapsed_transition_chunks = self.transition_chunks
|
|
|
|
def _current_alpha(self) -> float:
|
|
if self.transition_chunks <= 0:
|
|
return 1.0
|
|
return min(1.0, self.elapsed_transition_chunks / self.transition_chunks)
|
|
|
|
|
|
class _PendingPromptUpdate(TypedDict):
|
|
event_id: str
|
|
prompt: str
|
|
target_prompt_embeds: torch.Tensor
|
|
transition_chunks: int
|
|
|
|
|
|
class _PromptUpdateChunkMetadata(TypedDict):
|
|
started_event_ids: list[str]
|
|
active_event_ids: list[str]
|
|
completed_event_ids: list[str]
|