(
self,
key_states: torch.Tensor,
layer_idx: int,
)
| 115 | return new_cache |
| 116 | |
| 117 | def update_router_kcache( |
| 118 | self, |
| 119 | key_states: torch.Tensor, |
| 120 | layer_idx: int, |
| 121 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 122 | |
| 123 | # Update the cache |
| 124 | if key_states is not None: |
| 125 | if len(self.router_key_cache) <= layer_idx: |
| 126 | # There may be skipped layers, fill them with empty lists |
| 127 | for _ in range(len(self.router_key_cache), layer_idx): |
| 128 | self.router_key_cache.append(torch.tensor([])) |
| 129 | self.router_key_cache.append(key_states) |
| 130 | elif ( |
| 131 | not self.router_key_cache[layer_idx].numel() # prefers not t.numel() to len(t) == 0 to export the model |
| 132 | ): # fills previously skipped layers; checking for tensor causes errors |
| 133 | self.router_key_cache[layer_idx] = key_states |
| 134 | else: |
| 135 | self.router_key_cache[layer_idx] = torch.cat([self.router_key_cache[layer_idx], key_states], dim=-2) |
| 136 | |
| 137 | return self.router_key_cache[layer_idx] |
| 138 | |
| 139 | class CustomDynamicCacheOnCPU(CustomDynamicCache): |
| 140 | def __init__(self, _distributed_cache_data=None): |
no outgoing calls
no test coverage detected