52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
# 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)
|