| 89 | self.name = "report_doc_id" |
| 90 | |
| 91 | class CustomDynamicCacheOnCPU(CustomDynamicCache): |
| 92 | def __init__(self, _distributed_cache_data=None): |
| 93 | super().__init__(_distributed_cache_data) |
| 94 | |
| 95 | # def record_kwargs(self, layer_idx, kwargs): |
| 96 | # d = {} |
| 97 | # for k, v in kwargs.items(): |
| 98 | # if v is not None and torch.is_tensor(v): |
| 99 | # d[k] = v.cpu() if v.is_cuda else v.clone() |
| 100 | # else: |
| 101 | # d[k] = v |
| 102 | # super().record_kwargs(layer_idx, kwargs) |
| 103 | |
| 104 | def update( |
| 105 | self, |
| 106 | key_states: torch.Tensor, |
| 107 | value_states: torch.Tensor, |
| 108 | layer_idx: int, |
| 109 | cache_kwargs=None, |
| 110 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 111 | # if key_states is not None and torch.is_tensor(key_states) and key_states.is_cuda: |
| 112 | # key_states = key_states.cpu() |
| 113 | if value_states is not None and torch.is_tensor(value_states) and value_states.is_cuda: |
| 114 | value_states = value_states.cpu() |
| 115 | return super().update(key_states, value_states, layer_idx, cache_kwargs) |
| 116 | |
| 117 | @dataclass |
| 118 | class GenerateRequest(CmdBase): |
nothing calls this directly
no outgoing calls
no test coverage detected