Captures ``encoder_hidden_states.data_ptr()`` before ``txt_norm`` and writes it to shared state for the block hooks to read.
| 60 | |
| 61 | |
| 62 | class TextKVCacheTransformerHook(ModelHook): |
| 63 | """Captures ``encoder_hidden_states.data_ptr()`` before ``txt_norm`` |
| 64 | and writes it to shared state for the block hooks to read.""" |
| 65 | |
| 66 | _is_stateful = True |
| 67 | |
| 68 | def __init__(self, state_manager: StateManager): |
| 69 | super().__init__() |
| 70 | self.state_manager = state_manager |
| 71 | |
| 72 | def new_forward(self, module: torch.nn.Module, *args, **kwargs): |
| 73 | if self.state_manager._current_context is None: |
| 74 | self.state_manager.set_context("inference") |
| 75 | |
| 76 | encoder_hidden_states = kwargs.get("encoder_hidden_states") |
| 77 | if encoder_hidden_states is not None: |
| 78 | state: TextKVCacheState = self.state_manager.get_state() |
| 79 | state.key = encoder_hidden_states.data_ptr() |
| 80 | return self.fn_ref.original_forward(*args, **kwargs) |
| 81 | |
| 82 | def reset_state(self, module: torch.nn.Module): |
| 83 | self.state_manager.reset() |
| 84 | return module |
| 85 | |
| 86 | |
| 87 | class TextKVCacheBlockHook(ModelHook): |
no outgoing calls
no test coverage detected
searching dependent graphs…