[Perf][DreamZero]dreamzero supports CUDA graph, torch.compile, and DiT Caching. (#4213)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>

vllm_omni/diffusion/cache/stepcache/state.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Mutable state for step-level DiT velocity caching."""
+
+from __future__ import annotations
+
+import torch
+
+from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig
+
+
+class StepCacheState:
+    """Per-generation mutable state (skip countdown)."""
+
+    def __init__(self, config: StepCacheConfig) -> None:
+        self.config = config
+        self.skip_countdown = 0
+
+    def reset(self) -> None:
+        self.skip_countdown = 0
+
+    def should_run_step(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> bool:
+        """Return True when the DiT forward should execute at this scheduler step."""
+        if not self.config.enabled:
+            return True
+
+        if len(prev_predictions) < self.config.min_history_steps:
+            return True
+
+        if self.skip_countdown > 1:
+            self.skip_countdown -= 1
+            return False
+        if self.skip_countdown == 1:
+            self.skip_countdown = 0
+            return True
+
+        v_last = prev_predictions[-1][0].flatten(1).float()
+        v_prev = prev_predictions[-2][0].flatten(1).float()
+        sim = torch.nn.functional.cosine_similarity(v_last, v_prev, dim=1).mean()
+
+        for threshold, countdown in zip(self.config.sim_thresholds, self.config.skip_countdowns):
+            if sim > threshold:
+                self.skip_countdown = countdown
+                return False
+
+        return True
+
+    def trim_history(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> None:
+        while len(prev_predictions) > self.config.max_history:
+            prev_predictions.pop(0)