| from __future__ import annotations |
|
|
| from typing import Optional |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def _split_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor: |
| batch_size, seq_len, channels = x.shape |
| return x.reshape(batch_size, seq_len, num_heads, channels // num_heads).permute(0, 2, 1, 3).contiguous() |
|
|
|
|
| class SVDFasterCacheAttnProcessor: |
| ELIGIBLE_SEQ_LENS = {180, 720, 2880} |
|
|
| def __init__( |
| self, |
| *, |
| runtime_state: dict, |
| layer_idx: int, |
| first_layers_fp: int = 2, |
| ) -> None: |
| self.runtime_state = runtime_state |
| self.layer_idx = layer_idx |
| self.first_layers_fp = first_layers_fp |
|
|
| def __call__( |
| self, |
| attn, |
| hidden_states: torch.Tensor, |
| encoder_hidden_states: Optional[torch.Tensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| **kwargs, |
| ) -> torch.Tensor: |
| del kwargs |
| kv_source = encoder_hidden_states if encoder_hidden_states is not None else hidden_states |
| batch_size, query_len, _ = hidden_states.shape |
| key_len = kv_source.shape[1] |
| if not self._is_eligible(query_len=query_len, key_len=key_len): |
| return self._dense_forward( |
| attn, |
| hidden_states, |
| kv_source=kv_source, |
| attention_mask=attention_mask, |
| ) |
|
|
| history = self.runtime_state.setdefault("block_histories", {}).setdefault(self.layer_idx, []) |
| step_idx = int(self.runtime_state.get("current_step_idx", -1)) |
| start_step = int(self.runtime_state.get("resolved_start_step", 0)) |
| block_interval = max(int(self.runtime_state.get("block_interval", 3)), 1) |
| is_anchor_step = step_idx < start_step or ((step_idx - start_step) % block_interval == 0) |
| if ( |
| not is_anchor_step |
| and len(history) >= 2 |
| and history[-1].shape == hidden_states.shape |
| and history[-2].shape == hidden_states.shape |
| ): |
| latest = history[-1] |
| previous = history[-2] |
| return latest + (latest - previous) * float(self.runtime_state.get("block_alpha", 0.3)) |
|
|
| out = self._dense_forward( |
| attn, |
| hidden_states, |
| kv_source=kv_source, |
| attention_mask=attention_mask, |
| ) |
| history.append(out.detach().clone()) |
| if len(history) > 2: |
| del history[:-2] |
| return out |
|
|
| def _is_eligible(self, *, query_len: int, key_len: int) -> bool: |
| if self.layer_idx < self.first_layers_fp: |
| return False |
| if query_len != key_len: |
| return False |
| if query_len not in self.ELIGIBLE_SEQ_LENS: |
| return False |
| return True |
|
|
| @staticmethod |
| def _dense_forward( |
| attn, |
| hidden_states: torch.Tensor, |
| *, |
| kv_source: torch.Tensor, |
| attention_mask: Optional[torch.Tensor], |
| ) -> torch.Tensor: |
| num_heads = attn.heads |
| q = _split_heads(attn.to_q(hidden_states), num_heads) |
| k = _split_heads(attn.to_k(kv_source), num_heads) |
| v = _split_heads(attn.to_v(kv_source), num_heads) |
| hidden_states = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, dropout_p=0.0) |
| hidden_states = hidden_states.permute(0, 2, 1, 3).contiguous().reshape(hidden_states.shape[0], -1, attn.heads * q.shape[-1]) |
| hidden_states = attn.to_out[0](hidden_states) |
| hidden_states = attn.to_out[1](hidden_states) |
| return hidden_states |
|
|
