| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
|
|
| from .hooks import BaseState, HookRegistry, ModelHook, StateManager |
|
|
|
|
| _TEXT_KV_CACHE_TRANSFORMER_HOOK = "text_kv_cache_transformer" |
| _TEXT_KV_CACHE_BLOCK_HOOK = "text_kv_cache_block" |
|
|
|
|
| @dataclass |
| class TextKVCacheConfig: |
| """Enable exact (lossless) text K/V caching for transformer models. |
| |
| Pre-computes per-block text key and value projections once before the denoising loop and reuses them across all |
| steps. Positive and negative prompts are distinguished via a stable cache key captured by a transformer-level hook |
| before any intermediate tensor allocations. |
| """ |
|
|
| pass |
|
|
|
|
| class TextKVCacheState(BaseState): |
| """Shared state between the transformer-level and block-level hooks. |
| |
| The transformer hook writes the stable ``encoder_hidden_states`` ``data_ptr()`` (captured *before* ``txt_norm``) so |
| that block hooks can use it as a reliable cache key across denoising steps. |
| """ |
|
|
| def __init__(self): |
| self.key: int | None = None |
|
|
| def reset(self): |
| self.key = None |
|
|
|
|
| class TextKVCacheBlockState(BaseState): |
| """Per-block state holding cached text key/value projections.""" |
|
|
| def __init__(self): |
| self.kv_cache: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} |
|
|
| def reset(self): |
| self.kv_cache.clear() |
|
|
|
|
| class TextKVCacheTransformerHook(ModelHook): |
| """Captures ``encoder_hidden_states.data_ptr()`` before ``txt_norm`` |
| and writes it to shared state for the block hooks to read.""" |
|
|
| _is_stateful = True |
|
|
| def __init__(self, state_manager: StateManager): |
| super().__init__() |
| self.state_manager = state_manager |
|
|
| def new_forward(self, module: torch.nn.Module, *args, **kwargs): |
| if self.state_manager._current_context is None: |
| self.state_manager.set_context("inference") |
|
|
| encoder_hidden_states = kwargs.get("encoder_hidden_states") |
| if encoder_hidden_states is not None: |
| state: TextKVCacheState = self.state_manager.get_state() |
| state.key = encoder_hidden_states.data_ptr() |
| return self.fn_ref.original_forward(*args, **kwargs) |
|
|
| def reset_state(self, module: torch.nn.Module): |
| self.state_manager.reset() |
| return module |
|
|
|
|
| class TextKVCacheBlockHook(ModelHook): |
| """Caches ``(txt_key, txt_value)`` per block per unique prompt using |
| the stable cache key from the shared state.""" |
|
|
| _is_stateful = True |
|
|
| def __init__(self, state_manager: StateManager, block_state_manager: StateManager): |
| super().__init__() |
| self.state_manager = state_manager |
| self.block_state_manager = block_state_manager |
|
|
| def new_forward(self, module: torch.nn.Module, *args, **kwargs): |
| from ..models.transformers.transformer_nucleusmoe_image import _apply_rotary_emb_nucleus |
|
|
| if self.state_manager._current_context is None: |
| self.state_manager.set_context("inference") |
|
|
| if self.block_state_manager._current_context is None: |
| self.block_state_manager.set_context("inference") |
|
|
| if "encoder_hidden_states" in kwargs: |
| encoder_hidden_states = kwargs["encoder_hidden_states"] |
| else: |
| encoder_hidden_states = args[1] |
|
|
| if "image_rotary_emb" in kwargs: |
| image_rotary_emb = kwargs["image_rotary_emb"] |
| elif len(args) > 3: |
| image_rotary_emb = args[3] |
| else: |
| image_rotary_emb = None |
|
|
| state: TextKVCacheState = self.state_manager.get_state() |
| cache_key = state.key |
|
|
| block_state: TextKVCacheBlockState = self.block_state_manager.get_state() |
|
|
| if cache_key not in block_state.kv_cache: |
| context = module.encoder_proj(encoder_hidden_states) |
|
|
| attn = module.attn |
| head_dim = attn.inner_dim // attn.heads |
| num_kv_heads = attn.inner_kv_dim // head_dim |
|
|
| txt_key = attn.add_k_proj(context).unflatten(-1, (num_kv_heads, -1)) |
| txt_value = attn.add_v_proj(context).unflatten(-1, (num_kv_heads, -1)) |
|
|
| if attn.norm_added_k is not None: |
| txt_key = attn.norm_added_k(txt_key) |
|
|
| if image_rotary_emb is not None: |
| _, txt_freqs = image_rotary_emb |
| txt_key = _apply_rotary_emb_nucleus(txt_key, txt_freqs, use_real=False) |
|
|
| block_state.kv_cache[cache_key] = (txt_key, txt_value) |
|
|
| txt_key, txt_value = block_state.kv_cache[cache_key] |
|
|
| attn_kwargs = kwargs.get("attention_kwargs") or {} |
| attn_kwargs["cached_txt_key"] = txt_key |
| attn_kwargs["cached_txt_value"] = txt_value |
| kwargs["attention_kwargs"] = attn_kwargs |
|
|
| return self.fn_ref.original_forward(*args, **kwargs) |
|
|
| def reset_state(self, module: torch.nn.Module): |
| self.block_state_manager.reset() |
| return module |
|
|
|
|
| def apply_text_kv_cache(module: torch.nn.Module, config: TextKVCacheConfig) -> None: |
| from ..models.transformers.transformer_nucleusmoe_image import NucleusMoEImageTransformerBlock |
|
|
| HookRegistry.check_if_exists_or_initialize(module) |
|
|
| state_manager = StateManager(TextKVCacheState) |
|
|
| transformer_hook = TextKVCacheTransformerHook(state_manager) |
| registry = HookRegistry.check_if_exists_or_initialize(module) |
| registry.register_hook(transformer_hook, _TEXT_KV_CACHE_TRANSFORMER_HOOK) |
|
|
| for _, submodule in module.named_modules(): |
| if isinstance(submodule, NucleusMoEImageTransformerBlock): |
| block_state_manager = StateManager(TextKVCacheBlockState) |
| hook = TextKVCacheBlockHook(state_manager, block_state_manager) |
| block_registry = HookRegistry.check_if_exists_or_initialize(submodule) |
| block_registry.register_hook(hook, _TEXT_KV_CACHE_BLOCK_HOOK) |
|
|
