| 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): |