[Frontend] Add midway prompt update for streaming video generation (#4652)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

vllm_omni/diffusion/worker/input_batch.py

@@ -17,6 +17,7 @@ from dataclasses import dataclass, field
 import numpy as np
 import torch
 
+from vllm_omni.diffusion.prompt_update import prompt_update_versions
 from vllm_omni.diffusion.worker.utils import StepRequestState
 
 
@@ -535,12 +536,17 @@ def _same_composition(
     cached_batch: InputBatch | None,
     request_ids: list[str],
     idx_mapping_np: np.ndarray,
+    states: Sequence[StepRequestState],
 ) -> bool:
     if cached_batch is None:
         return False
     if cached_batch.request_ids != request_ids:
         return False
-    return np.array_equal(cached_batch.idx_mapping_np, idx_mapping_np)
+    if not np.array_equal(cached_batch.idx_mapping_np, idx_mapping_np):
+        return False
+    # Midway prompt updates (typically for video generation) can change prompt_embeds without changing ids/mapping.
+    # In this case, each request state manages the embedding's "version". Use it to determine if cache is still valid.
+    return cached_batch._prompt_update_versions == prompt_update_versions(states)
 
 
 def _scatter_batch_tensor_by_mapping(
@@ -615,6 +621,10 @@ class InputBatch:
     negative_txt_seq_lens: list[int] | None = None
     states: Sequence[StepRequestState] = field(default_factory=tuple)
 
+    # For midway prompt updates (typically for video generation) that changes embeddings without changing ids,
+    # Keep a snapshot of the current per-request versions of prompt embeddings for later runtime comparison
+    _prompt_update_versions: tuple[int, ...] = ()
+
     def __post_init__(self) -> None:
         if len(self.request_ids) != int(self.idx_mapping.numel()):
             raise ValueError("`request_ids` and `idx_mapping` must have the same length.")
@@ -653,6 +663,7 @@ class InputBatch:
         self.img_shapes = _prepare_img_shapes(states)
         self.txt_seq_lens = _prepare_seq_lens(states, "txt_seq_lens")
         self.negative_txt_seq_lens = _prepare_seq_lens(states, "negative_txt_seq_lens")
+        self._prompt_update_versions = prompt_update_versions(states)
 
     def _repack_dynamic_fields(
         self,
@@ -693,7 +704,7 @@ class InputBatch:
         selected_states, idx_mapping, idx_mapping_np = _select_states(states, idx_mapping)
         request_ids = _prepare_request_ids(selected_states)
 
-        if _same_composition(cached_batch, request_ids, idx_mapping_np):
+        if _same_composition(cached_batch, request_ids, idx_mapping_np, selected_states):
             assert cached_batch is not None
             cached_batch._repack_dynamic_fields(selected_states)
             return cached_batch
@@ -733,6 +744,7 @@ class InputBatch:
                 "negative_txt_seq_lens",
             ),
             states=tuple(selected_states),
+            _prompt_update_versions=prompt_update_versions(selected_states),
         )