# Copyright 2025 The HuggingFace Team. All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

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)