(module: torch.nn.Module, config: TextKVCacheConfig)
| 155 | |
| 156 | |
| 157 | def apply_text_kv_cache(module: torch.nn.Module, config: TextKVCacheConfig) -> None: |
| 158 | from ..models.transformers.transformer_nucleusmoe_image import NucleusMoEImageTransformerBlock |
| 159 | |
| 160 | HookRegistry.check_if_exists_or_initialize(module) |
| 161 | |
| 162 | state_manager = StateManager(TextKVCacheState) |
| 163 | |
| 164 | transformer_hook = TextKVCacheTransformerHook(state_manager) |
| 165 | registry = HookRegistry.check_if_exists_or_initialize(module) |
| 166 | registry.register_hook(transformer_hook, _TEXT_KV_CACHE_TRANSFORMER_HOOK) |
| 167 | |
| 168 | for _, submodule in module.named_modules(): |
| 169 | if isinstance(submodule, NucleusMoEImageTransformerBlock): |
| 170 | block_state_manager = StateManager(TextKVCacheBlockState) |
| 171 | hook = TextKVCacheBlockHook(state_manager, block_state_manager) |
| 172 | block_registry = HookRegistry.check_if_exists_or_initialize(submodule) |
| 173 | block_registry.register_hook(hook, _TEXT_KV_CACHE_BLOCK_HOOK) |
no test coverage detected
searching dependent graphs…