自定义动态缓存类 扩展标准的DynamicCache，添加额外的元数据存储和查询功能
| 12 | |
| 13 | |
| 14 | class CustomDynamicCache(DynamicCache): |
| 15 | """ |
| 16 | 自定义动态缓存类 |
| 17 | 扩展标准的DynamicCache，添加额外的元数据存储和查询功能 |
| 18 | """ |
| 19 | def __init__(self, _distributed_cache_data=None): |
| 20 | super().__init__(_distributed_cache_data) |
| 21 | self.cache_kwargs = {} |
| 22 | self.group_cache = {} |
| 23 | self.meta = {} |
| 24 | self.router_key_cache = [] |
| 25 | |
| 26 | def clear_kvcache(self): |
| 27 | self.key_cache = [] |
| 28 | self.value_cache = [] |
| 29 | |
| 30 | def record_kwargs(self, layer_idx, kwargs): |
| 31 | """ |
| 32 | 记录层的元数据信息 |
| 33 | |
| 34 | Args: |
| 35 | layer_idx: 层索引 |
| 36 | kwargs: 包含路由层信息的字典 |
| 37 | """ |
| 38 | if layer_idx in self.cache_kwargs: |
| 39 | self.cache_kwargs[layer_idx].update(kwargs) |
| 40 | else: |
| 41 | self.cache_kwargs[layer_idx] = kwargs |
| 42 | |
| 43 | def get_layer_length(self): |
| 44 | return len(self.cache_kwargs) |
| 45 | |
| 46 | def get_kvcache(self, layer_idx): |
| 47 | """ |
| 48 | 获取指定层的KV缓存 |
| 49 | |
| 50 | Args: |
| 51 | layer_idx: 层索引 |
| 52 | |
| 53 | Returns: |
| 54 | Tuple[torch.Tensor, torch.Tensor]: (key_cache, value_cache) |
| 55 | """ |
| 56 | key_cache = self.key_cache[layer_idx] |
| 57 | value_cache = self.value_cache[layer_idx] |
| 58 | return key_cache, value_cache |
| 59 | |
| 60 | def get_router_kcache(self, layer_idx): |
| 61 | if layer_idx < len(self.router_key_cache): |
| 62 | return self.router_key_cache[layer_idx] |
| 63 | else: |
| 64 | return None |
| 65 | |
| 66 | def clear_query(self): |
| 67 | """ |
| 68 | 清理查询相关的临时数据 |
| 69 | 移除查询过程中产生的临时数据，保持缓存清洁 |
| 70 | """ |
| 71 | for k, v in self.cache_kwargs.items(): |
no outgoing calls
no test coverage detected