Caches ``(txt_key, txt_value)`` per block per unique prompt using the stable cache key from the shared state.
| 85 | |
| 86 | |
| 87 | class TextKVCacheBlockHook(ModelHook): |
| 88 | """Caches ``(txt_key, txt_value)`` per block per unique prompt using |
| 89 | the stable cache key from the shared state.""" |
| 90 | |
| 91 | _is_stateful = True |
| 92 | |
| 93 | def __init__(self, state_manager: StateManager, block_state_manager: StateManager): |
| 94 | super().__init__() |
| 95 | self.state_manager = state_manager |
| 96 | self.block_state_manager = block_state_manager |
| 97 | |
| 98 | def new_forward(self, module: torch.nn.Module, *args, **kwargs): |
| 99 | from ..models.transformers.transformer_nucleusmoe_image import _apply_rotary_emb_nucleus |
| 100 | |
| 101 | if self.state_manager._current_context is None: |
| 102 | self.state_manager.set_context("inference") |
| 103 | |
| 104 | if self.block_state_manager._current_context is None: |
| 105 | self.block_state_manager.set_context("inference") |
| 106 | |
| 107 | if "encoder_hidden_states" in kwargs: |
| 108 | encoder_hidden_states = kwargs["encoder_hidden_states"] |
| 109 | else: |
| 110 | encoder_hidden_states = args[1] |
| 111 | |
| 112 | if "image_rotary_emb" in kwargs: |
| 113 | image_rotary_emb = kwargs["image_rotary_emb"] |
| 114 | elif len(args) > 3: |
| 115 | image_rotary_emb = args[3] |
| 116 | else: |
| 117 | image_rotary_emb = None |
| 118 | |
| 119 | state: TextKVCacheState = self.state_manager.get_state() |
| 120 | cache_key = state.key |
| 121 | |
| 122 | block_state: TextKVCacheBlockState = self.block_state_manager.get_state() |
| 123 | |
| 124 | if cache_key not in block_state.kv_cache: |
| 125 | context = module.encoder_proj(encoder_hidden_states) |
| 126 | |
| 127 | attn = module.attn |
| 128 | head_dim = attn.inner_dim // attn.heads |
| 129 | num_kv_heads = attn.inner_kv_dim // head_dim |
| 130 | |
| 131 | txt_key = attn.add_k_proj(context).unflatten(-1, (num_kv_heads, -1)) |
| 132 | txt_value = attn.add_v_proj(context).unflatten(-1, (num_kv_heads, -1)) |
| 133 | |
| 134 | if attn.norm_added_k is not None: |
| 135 | txt_key = attn.norm_added_k(txt_key) |
| 136 | |
| 137 | if image_rotary_emb is not None: |
| 138 | _, txt_freqs = image_rotary_emb |
| 139 | txt_key = _apply_rotary_emb_nucleus(txt_key, txt_freqs, use_real=False) |
| 140 | |
| 141 | block_state.kv_cache[cache_key] = (txt_key, txt_value) |
| 142 | |
| 143 | txt_key, txt_value = block_state.kv_cache[cache_key] |
| 144 |
no outgoing calls
no test coverage detected
searching dependent graphs…