Aluode/PerceptionLabPortable
0
1from abc import ABC, abstractmethod2from collections.abc import Iterable3from typing import Any, Optional4 5import torch6 7from .configuration_utils import PretrainedConfig8from .utils import (9 is_hqq_available,10 is_quanto_greater,11 is_torch_greater_or_equal,12 is_torchdynamo_compiling,13 logging,14)15 16 17if is_hqq_available():18 from hqq.core.quantize import Quantizer as HQQQuantizer19 20_is_torch_greater_or_equal_than_2_7 = is_torch_greater_or_equal("2.7", accept_dev=True)21 22 23logger = logging.get_logger(__name__)24 25 26class CacheLayerMixin(ABC):27 """Base, abstract class for a single layer's cache."""28 29 is_compileable = False30 31 def __init__(self):32 self.keys: Optional[torch.Tensor] = None33 self.values: Optional[torch.Tensor] = None34 self.is_initialized = False35 36 def __repr__(self):37 return f"{self.__class__.__name__}"38 39 @abstractmethod40 def lazy_initialization(self, key_states: torch.Tensor): ...41 42 @abstractmethod43 def update(44 self, key_states: torch.Tensor, value_states: torch.Tensor, cache_kwargs: Optional[dict[str, Any]] = None45 ) -> tuple[torch.Tensor, torch.Tensor]: ...46 47 @abstractmethod48 def get_mask_sizes(self, cache_position: torch.Tensor) -> tuple[int, int]: ...49 50 @abstractmethod51 def get_seq_length(self) -> int: ...52 53 @abstractmethod54 def get_max_cache_shape(self) -> int: ...55 56 def offload(self):57 """Offload this layer's data to CPU device."""58 if self.is_initialized:59 self.keys = self.keys.to("cpu", non_blocking=True)60 self.values = self.values.to("cpu", non_blocking=True)61 62 def prefetch(self):63 """In case of layer offloading, this allows to move the data back to the layer's device ahead of time."""64 if self.is_initialized and self.keys.device != self.device:65 self.keys = self.keys.to(self.device, non_blocking=True)66 self.values = self.values.to(self.device, non_blocking=True)67 68 def reset(self) -> None:69 """Resets the cache values while preserving the objects"""70 if self.is_initialized:71 self.keys.zero_()72 self.values.zero_()73 # This attribute is set on several Layers74 if hasattr(self, "cumulative_length"):75 self.cumulative_length = 076 77 def reorder_cache(self, beam_idx: torch.LongTensor) -> None:78 """Reorders this layer's cache for beam search."""79 if self.get_seq_length() > 0:80 self.keys = self.keys.index_select(0, beam_idx.to(self.keys.device))81 self.values = self.values.index_select(0, beam_idx.to(self.values.device))82 83 84class DynamicLayer(CacheLayerMixin):85 """86 A cache layer that grows dynamically as more tokens are generated. This is the default for generative models.87 It stores the key and value states as tensors of shape `[batch_size, num_heads, seq_len, head_dim]`.88 """89 90 is_sliding = False91 92 def lazy_initialization(self, key_states: torch.Tensor):93 self.dtype, self.device = key_states.dtype, key_states.device94 self.keys = torch.tensor([], dtype=self.dtype, device=self.device)95 self.values = torch.tensor([], dtype=self.dtype, device=self.device)96 self.is_initialized = True97 98 def update(99 self,100 key_states: torch.Tensor,101 value_states: torch.Tensor,102 cache_kwargs: Optional[dict[str, Any]] = None,103 ) -> tuple[torch.Tensor, torch.Tensor]:104 """105 Update the key and value caches in-place, and return the necessary keys and value states.106 107 Args:108 key_states (`torch.Tensor`): The new key states to cache.109 value_states (`torch.Tensor`): The new value states to cache.110 cache_kwargs (`dict[str, Any]`, *optional*): Additional arguments for the cache.111 112 Returns:113 tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states.114 """115 # Lazy initialization116 if not self.is_initialized:117 self.lazy_initialization(key_states)118 119 self.keys = torch.cat([self.keys, key_states], dim=-2)120 self.values = torch.cat([self.values, value_states], dim=-2)121 return self.keys, self.values122 123 def get_mask_sizes(self, cache_position: torch.Tensor) -> tuple[int, int]:124 """Return the length and offset of the cache, used to generate the mask"""125 kv_offset = 0126 query_length = cache_position.shape[0]127 kv_length = self.get_seq_length() + query_length128 return kv_length, kv_offset129 130 def get_seq_length(self) -> int:131 """Returns the sequence length of the cached states."""132 if not self.is_initialized or self.keys.numel() == 0:133 return 0134 return self.keys.shape[-2]135 136 def get_max_cache_shape(self) -> int:137 """Returns the maximum sequence length of the cache object. DynamicLayer does not have a maximum length."""138 return -1139 140 def crop(self, max_length: int) -> None:141 """142 Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be negative143 to remove `max_length` tokens.144 """145 if max_length < 0:146 max_length = self.get_seq_length() - abs(max_length)147 148 if self.get_seq_length() <= max_length:149 return150 151 self.keys = self.keys[..., :max_length, :]152 self.values = self.values[..., :max_length, :]153 154 def batch_repeat_interleave(self, repeats: int) -> None:155 """Repeat the cache `repeats` times in the batch dimension."""156 if self.get_seq_length() > 0:157 self.keys = self.keys.repeat_interleave(repeats, dim=0)158 self.values = self.values.repeat_interleave(repeats, dim=0)159 160 def batch_select_indices(self, indices: torch.Tensor) -> None:161 """Only keep the `indices` in the batch dimension of the cache."""162 if self.get_seq_length() > 0:163 self.keys = self.keys[indices, ...]164 self.values = self.values[indices, ...]165 166 167class DynamicSlidingWindowLayer(DynamicLayer):168 """169 A cache layer that grows dynamically as more tokens are generated, up until the sliding window size.170 It stores the key and value states as tensors of shape `[batch_size, num_heads, min(seq_len, sliding_window), head_dim]`.171 """172 173 is_sliding = True174 175 def __init__(self, sliding_window: int):176 super().__init__()177 self.sliding_window = sliding_window178 self.cumulative_length = 0179 180 def update(181 self,182 key_states: torch.Tensor,183 value_states: torch.Tensor,184 cache_kwargs: Optional[dict[str, Any]] = None,185 ) -> tuple[torch.Tensor, torch.Tensor]:186 """187 Update the key and value caches in-place, and return the necessary keys and value states.188 189 Args:190 key_states (`torch.Tensor`): The new key states to cache.191 value_states (`torch.Tensor`): The new value states to cache.192 cache_kwargs (`dict[str, Any]`, *optional*): Additional arguments for the cache.193 194 Returns:195 tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states.196 """197 # Lazy initialization198 if not self.is_initialized:199 self.lazy_initialization(key_states)200 201 self.cumulative_length += key_states.shape[-2]202 203 # Compute the full states204 full_key_states = torch.cat([self.keys, key_states], dim=-2)205 full_value_states = torch.cat([self.values, value_states], dim=-2)206 # Only cache the last `self.sliding_window - 1` tokens (or all of them if lower than that)207 self.keys = full_key_states[:, :, -self.sliding_window + 1 :, :]208 self.values = full_value_states[:, :, -self.sliding_window + 1 :, :]209 210 # Return the full states211 return full_key_states, full_value_states212 213 def get_mask_sizes(self, cache_position: torch.Tensor) -> tuple[int, int]:214 """Return the length and offset of the cache, used to generate the attention mask"""215 query_length = cache_position.shape[0]216 is_full = self.cumulative_length >= self.sliding_window217 218 kv_offset = max(self.cumulative_length - self.sliding_window + 1, 0)219 if is_full:220 kv_length = self.sliding_window - 1 + query_length221 else:222 kv_length = self.cumulative_length + query_length223 224 return kv_length, kv_offset225 226 def get_seq_length(self) -> int:227 """Returns the sequence length of the cached states."""228 return self.cumulative_length229 230 def get_max_cache_shape(self) -> int:231 """Return the maximum cache shape of the cache"""232 return self.sliding_window233 234 def crop(self, max_length: int) -> None:235 """236 Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be237 negative to remove `max_length` tokens.238 """239 if self.get_seq_length() >= self.sliding_window:240 raise ValueError(241 "Cannot `crop` a `DynamicSlidingWindowLayer` after it has seen more tokens than its"242 "sliding window (otherwise some states are lost)"243 )244 super().crop(max_length)245 self.cumulative_length = self.keys.shape[-2]246 247 248class StaticLayer(CacheLayerMixin):249 """250 A static cache layer that stores the key and value states as static tensors of shape `[batch_size, num_heads, max_cache_len), head_dim]`.251 It lazily allocates its full backing tensors, and then mutates them in-place. Built for `torch.compile` support.252 253 Args:254 max_cache_len (`int`):255 Maximum number of tokens that can be stored, used for tensor preallocation.256 """257 258 is_compileable = True259 is_sliding = False260 261 def __init__(self, max_cache_len: int):262 super().__init__()263 self.max_cache_len = max_cache_len264 265 def lazy_initialization(self, key_states: torch.Tensor):266 """267 Lazy initialization of the keys and values tensors. This allows to get all properties (dtype, device,268 num_heads in case of TP etc...) at runtime directly, which is extremely practical as it avoids moving269 devices, dtypes etc later on for each `update` (which could break the static dynamo addresses as well).270 271 If this is unwanted, one can call `early_initialization(...)` on the Cache directly, which will call this272 function ahead-of-time (this is required for `torch.export` for example). Note that for `compile`, as we273 internally don't compile the prefill, this is guaranteed to have been called already when compiling.274 If compiling the prefill as well, e.g. calling `model.compile(...)` before `generate` with a static cache,275 it is still supported in general, but without guarantees depending on the compilation options (e.g. cuda graphs,276 i.e. `mode="reduce-overhead"` is known to fail). But it will in general work correctly, and prefill should277 not be compiled anyway for performances!278 """279 self.max_batch_size, self.num_heads, _, self.head_dim = key_states.shape280 self.dtype, self.device = key_states.dtype, key_states.device281 282 self.keys = torch.zeros(283 (self.max_batch_size, self.num_heads, self.max_cache_len, self.head_dim),284 dtype=self.dtype,285 device=self.device,286 )287 self.values = torch.zeros(288 (self.max_batch_size, self.num_heads, self.max_cache_len, self.head_dim),289 dtype=self.dtype,290 device=self.device,291 )292 # Note: `mark_static_address` is used to tag the cache as a fixed data pointer, preventing compiled graph293 # breaks when updating the cache. However, it is not supported when tracing the graph, so we skip it in this case.294 # As prefill should never be compiled, this is not an issue and it will still be run (except when users compile295 # prefill explicitly, but this should be avoided!)296 if not is_torchdynamo_compiling():297 torch._dynamo.mark_static_address(self.keys)298 torch._dynamo.mark_static_address(self.values)299 300 self.is_initialized = True301 302 def update(303 self,304 key_states: torch.Tensor,305 value_states: torch.Tensor,306 cache_kwargs: Optional[dict[str, Any]] = None,307 ) -> tuple[torch.Tensor, torch.Tensor]:308 """309 Update the key and value caches in-place, and return the necessary keys and value states.310 311 Args:312 key_states (`torch.Tensor`): The new key states to cache.313 value_states (`torch.Tensor`): The new value states to cache.314 cache_kwargs (`dict[str, Any]`, *optional*): Additional arguments for the cache.315 316 Returns:317 tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states.318 """319 # Lazy initialization320 if not self.is_initialized:321 self.lazy_initialization(key_states)322 323 # Some old models give None for `cache_position` or even omit passing `cache_kwargs` when used as cross-attention,324 # in which case we should copy the whole Layer (key_states.shape[-2] == self.max_cache_len)325 cache_position = cache_kwargs.get("cache_position") if cache_kwargs is not None else None326 cache_position = (327 cache_position if cache_position is not None else torch.arange(key_states.shape[-2], device=self.device)328 )329 330 # Update the cache331 try:332 self.keys.index_copy_(2, cache_position, key_states)333 self.values.index_copy_(2, cache_position, value_states)334 except NotImplementedError:335 # Fallback for devices like MPS where index_copy_ might not be supported.336 self.keys[:, :, cache_position] = key_states337 self.values[:, :, cache_position] = value_states338 return self.keys, self.values339 340 def get_mask_sizes(self, cache_position: torch.Tensor) -> tuple[int, int]:341 """Return the length and offset of the cache, used to generate the attention mask"""342 kv_offset = 0343 kv_length = self.max_cache_len344 return kv_length, kv_offset345 346 def get_seq_length(self) -> int:347 """Returns the sequence length of the cached states."""348 # Occupied cache == any slot in the 3rd dim (sequence length) holds a non-zero value. To save on compute, let's349 # limit the check to the first batch member and head dimension.350 return (self.keys[0, 0].any(dim=-1)).sum() if self.is_initialized else 0351 352 def get_max_cache_shape(self) -> int:353 """Return the maximum cache shape of the cache"""354 return self.max_cache_len355 356 357class StaticSlidingWindowLayer(StaticLayer):358 """359 A static cache layer that stores the key and value states as static tensors of shape360 `[batch_size, num_heads, min(max_cache_len, sliding_window), head_dim]`. It lazily allocates its full backing361 tensors, and then mutates them in-place. Built for `torch.compile` support.362 363 Args:364 max_cache_len (`int`):365 Maximum number of tokens that can be stored, used for tensor preallocation.366 sliding_window (`int`):367 The size of the sliding window.368 """369 370 is_sliding = True371 372 def __init__(self, max_cache_len: int, sliding_window: int):373 effective_max_cache_len = min(sliding_window, max_cache_len)374 super().__init__(max_cache_len=effective_max_cache_len)375 self.cumulative_length = 0376 377 def update(378 self,379 key_states: torch.Tensor,380 value_states: torch.Tensor,381 cache_kwargs: Optional[dict[str, Any]] = None,382 ) -> tuple[torch.Tensor, torch.Tensor]:383 """384 Update the key and value caches in-place, and return the necessary keys and value states.385 386 Args:387 key_states (`torch.Tensor`): The new key states to cache.388 value_states (`torch.Tensor`): The new value states to cache.389 cache_kwargs (`dict[str, Any]`, *optional*): Additional arguments for the cache.390 391 Returns:392 tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states.393 """394 # Lazy initialization395 if not self.is_initialized:396 self.lazy_initialization(key_states)397 398 # Some old models give None for `cache_position` or even omit passing `cache_kwargs` when used as cross-attention,399 # in which case we should copy the whole Layer (key_states.shape[-2] == self.max_cache_len)400 cache_position = cache_kwargs.get("cache_position") if cache_kwargs is not None else None401 cache_position = (402 cache_position if cache_position is not None else torch.arange(key_states.shape[-2], device=self.device)403 )404 405 cumulative_length = self.cumulative_length406 is_full = cumulative_length >= self.max_cache_len407 # Update it now that we saved the value above408 self.cumulative_length += key_states.shape[-2]409 410 if is_full:411 # In general, we should use a much simpler `cat` here as well, independently of the states size. However,412 # dynamo is currently bugged when doing it - see https://github.com/pytorch/pytorch/issues/159855 for more details413 if key_states.shape[-2] == 1:414 # Roll all values to the left by 1 position415 new_keys = self.keys.roll(-1, dims=-2)416 new_values = self.values.roll(-1, dims=-2)417 # Overwrite the last position with new states418 # (note: very important to use a tensor to index here, see https://github.com/pytorch/pytorch/issues/159855)419 index = torch.tensor([-1], dtype=int, device=self.device)420 new_keys[:, :, index] = key_states421 new_values[:, :, index] = value_states422 423 # Copy back into `self` (do not just assign again) in order to keep the static dynamo address424 self.keys.copy_(new_keys)425 self.values.copy_(new_values)426 # Very important to return the `self` tensors here, as they have the static dynamo address427 return self.keys, self.values428 # Already full but using more than 1 new token (e.g. prefill caching, chat continuation, etc...)429 else:430 full_key_states = torch.cat((self.keys[:, :, 1:, :], key_states), dim=-2)431 full_value_states = torch.cat((self.values[:, :, 1:, :], value_states), dim=-2)432 # Not yet full, but becoming full on this update433 elif cumulative_length + key_states.shape[2] > self.max_cache_len:434 # Fast prefill path, no need to cat() in this case, as the cache is currently empty435 if cumulative_length == 0:436 full_key_states = key_states437 full_value_states = value_states438 else:439 full_key_states = torch.cat((self.keys[:, :, :cumulative_length, :], key_states), dim=-2)440 full_value_states = torch.cat((self.values[:, :, :cumulative_length, :], value_states), dim=-2)441 else:442 try:443 self.keys.index_copy_(2, cache_position, key_states)444 self.values.index_copy_(2, cache_position, value_states)445 except NotImplementedError:446 self.keys[:, :, cache_position] = key_states447 self.values[:, :, cache_position] = value_states448 449 # Very important to return the `self` tensors here, as they have the static dynamo address450 return self.keys, self.values451 452 # We only cache the last `sliding_window` tokens453 self.keys.copy_(full_key_states[:, :, -self.max_cache_len :, :])454 self.values.copy_(full_value_states[:, :, -self.max_cache_len :, :])455 # we should return the whole states instead of `self.keys/values` here, as otherwise we lose some context456 return full_key_states, full_value_states457 458 def get_mask_sizes(self, cache_position: torch.Tensor) -> tuple[int, int]:459 """Return the length and offset of the cache, used to generate the attention mask"""460 query_length = cache_position.shape[0]461 sliding_window = self.max_cache_len462 is_full = self.cumulative_length >= self.max_cache_len463 464 kv_offset = max(self.cumulative_length - sliding_window + 1, 0)465 # The cache is already full466 if is_full:467 kv_length = sliding_window + query_length - 1468 # Not yet full, but becoming full on this update469 elif self.cumulative_length + query_length > sliding_window:470 kv_length = self.cumulative_length + query_length471 # Here the Cache is still smaller than the local size, but we return the local size as it's static472 else:473 kv_length = sliding_window474 475 return kv_length, kv_offset476 477 def get_seq_length(self) -> int:478 """Returns the sequence length of the cached states."""479 return self.cumulative_length480 481 482class QuantizedLayer(DynamicLayer):483 """484 A quantized layer similar to what is described in the [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache paper](https://huggingface.co/papers/2402.02750).485 It allows the model to generate longer sequence length without allocating too much memory for the key and value caches by486 applying quantization.487 488 The cache has two types of storage, one for original precision and one for the quantized cache. A `residual length`489 is set as a maximum capacity for the original precision cache. When the length goes beyond maximum capacity, the original490 precision cache is discarded and moved into the quantized cache. The quantization is done per-channel with a set `q_group_size`491 for both Keys and Values, in contrast to what was described in the paper.492 """493 494 def __init__(495 self,496 nbits: int = 4,497 axis_key: int = 0,498 axis_value: int = 0,499 q_group_size: int = 64,500 residual_length: int = 128,501 ):502 super().__init__()503 self.nbits = nbits504 self.axis_key = axis_key505 self.axis_value = axis_value506 self.q_group_size = q_group_size507 self.residual_length = residual_length508 self.cumulative_length = 0509 510 def update(511 self,512 key_states: torch.Tensor,513 value_states: torch.Tensor,514 cache_kwargs: Optional[dict[str, Any]] = None,515 ) -> tuple[torch.Tensor, torch.Tensor]:516 """517 Update the key and value caches in-place, and return the necessary keys and value states.518 519 Args:520 key_states (`torch.Tensor`): The new key states to cache.521 value_states (`torch.Tensor`): The new value states to cache.522 cache_kwargs (`dict[str, Any]`, *optional*): Additional arguments for the cache.523 524 Returns:525 tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states.526 """527 self.cumulative_length += key_states.shape[-2]528 529 # Lazy initialization530 if not self.is_initialized:531 self.lazy_initialization(key_states)532 self._quantized_keys = self._quantize(key_states.contiguous(), axis=self.axis_key)533 self._quantized_values = self._quantize(value_states.contiguous(), axis=self.axis_value)534 return key_states, value_states535 536 dequant_keys = self._dequantize(self._quantized_keys)537 dequant_values = self._dequantize(self._quantized_values)538 keys_to_return = torch.cat([dequant_keys, self.keys, key_states], dim=-2)539 values_to_return = torch.cat([dequant_values, self.values, value_states], dim=-2)540 if self.keys.dim() == 4 and self.keys.shape[-2] + 1 >= self.residual_length:541 self._quantized_keys = self._quantize(keys_to_return.contiguous(), axis=self.axis_key)542 self._quantized_values = self._quantize(values_to_return.contiguous(), axis=self.axis_value)543 self.keys = torch.tensor([], dtype=key_states.dtype, device=key_states.device)544 self.values = torch.tensor([], dtype=key_states.dtype, device=key_states.device)545 else:546 self.keys = torch.cat([self.keys, key_states], dim=-2)547 self.values = torch.cat([self.values, value_states], dim=-2)548 549 return keys_to_return, values_to_return550 551 @abstractmethod552 def _quantize(self, tensor, axis): ...553 554 @abstractmethod555 def _dequantize(self, q_tensor): ...556 557 def get_seq_length(self) -> int:558 """Returns the sequence length of the cached states."""559 return self.cumulative_length560 561 562class QuantoQuantizedLayer(QuantizedLayer):563 def __init__(564 self,565 nbits: int = 4,566 axis_key: int = 0,567 axis_value: int = 0,568 q_group_size: int = 64,569 residual_length: int = 128,570 ):571 super().__init__(572 nbits=nbits,573 axis_key=axis_key,574 axis_value=axis_value,575 q_group_size=q_group_size,576 residual_length=residual_length,577 )578 579 # We need to import quanto here to avoid circular imports due to optimum/quanto/models/transformers_models.py580 if is_quanto_greater("0.2.5", accept_dev=True):581 from optimum.quanto import MaxOptimizer, qint2, qint4582 else:583 raise ImportError(584 "You need optimum-quanto package version to be greater or equal than 0.2.5 to use `QuantoQuantizedCache`. "585 )586 587 if self.nbits not in [2, 4]:588 raise ValueError(f"`nbits` for `quanto` backend has to be one of [`2`, `4`] but got {self.nbits}")589 590 if self.axis_key not in [0, -1]:591 raise ValueError(f"`axis_key` for `quanto` backend has to be one of [`0`, `-1`] but got {self.axis_key}")592 593 if self.axis_value not in [0, -1]:594 raise ValueError(595 f"`axis_value` for `quanto` backend has to be one of [`0`, `-1`] but got {self.axis_value}"596 )597 598 self.qtype = qint4 if self.nbits == 4 else qint2599 self.optimizer = MaxOptimizer() # hardcode as it's the only one for per-channel quantization600 601 def _quantize(self, tensor, axis):602 from optimum.quanto import quantize_weight603 604 scale, zeropoint = self.optimizer(tensor, self.qtype, axis, self.q_group_size)605 qtensor = quantize_weight(tensor, self.qtype, axis, scale, zeropoint, self.q_group_size)606 return qtensor607 608 def _dequantize(self, qtensor):609 return qtensor.dequantize()610 611 612class HQQQuantizedLayer(QuantizedLayer):613 def __init__(614 self,615 nbits: int = 4,616 axis_key: int = 0,617 axis_value: int = 0,618 q_group_size: int = 64,619 residual_length: int = 128,620 ):621 super().__init__(622 nbits=nbits,623 axis_key=axis_key,624 axis_value=axis_value,625 q_group_size=q_group_size,626 residual_length=residual_length,627 )628 629 if not is_hqq_available():630 raise ImportError("You need to install `hqq` to use `HQQQuantizedLayer`")631 632 if self.nbits not in [1, 2, 3, 4, 8]:633 raise ValueError(634 f"`nbits` for `HQQ` backend has to be one of [`1`, `2`, `3`, `4`, `8`] but got {self.nbits}"635 )636 637 if self.axis_key not in [0, 1]:638 raise ValueError(f"`axis_key` for `HQQ` backend has to be one of [`0`, `1`] but got {self.axis_key}")639 640 if self.axis_value not in [0, 1]:641 raise ValueError(f"`axis_value` for `HQQ` backend has to be one of [`0`, `1`] but got {self.axis_value}")642 643 self.quantizer = HQQQuantizer644 645 def _quantize(self, tensor, axis):646 qtensor, meta = self.quantizer.quantize(647 tensor,648 axis=axis,649 device=self.keys.device,650 compute_dtype=self.keys.dtype,651 nbits=self.nbits,652 group_size=self.q_group_size,653 )654 meta["compute_dtype"] = self.keys.dtype655 self.quantizer.cuda(qtensor, meta=meta, device=self.keys.device) # Move to device and cast to dtype656 meta["scale"] = meta["scale"].to(qtensor.device)657 meta["zero"] = meta["zero"].to(qtensor.device)658 return qtensor, meta659 660 def _dequantize(self, qtensor):661 quant_tensor, meta = qtensor662 tensor = self.quantizer.dequantize(quant_tensor, meta)663 return tensor664 665 666class Cache:667 """668 A `Cache` is mostly a list of `CacheLayerMixin` objects, one per model layer. It serves as a container for669 the Cache of each layer.670 671 Args:672 layers (`Optional`, *optional*):673 A list of pre-created `CacheLayerMixin`. If omitted (`None`), then `layer_class_to_replicate` will674 be used.675 layer_class_to_replicate (`type[CacheLayerMixin]`, *optional*):676 Only used if `layers` is omitted (`None`), in which case it will be used as the base class for each layer,677 and the layers will be added lazily as soon as `update` is called with a `layer_idx` greater than the current678 list of layers.679 offloading (`bool`, *optional*, defaults to `False`):680 Whether to perform offloading of the layers to `cpu`, to save GPU memory.681 offload_only_non_sliding (`bool`, *optional*, defaults to `True`):682 If `offloading` is `True`, this further decides if only the non-sliding layers will be offloaded (because683 usually the sliding layers are small in size, so there is no need to offload them, and skipping it is faster).684 """685 686 def __init__(687 self,688 layers: Optional[list[CacheLayerMixin]] = None,689 layer_class_to_replicate: Optional[type[CacheLayerMixin]] = None,690 offloading: bool = False,691 offload_only_non_sliding: bool = True,692 ):693 if layers is not None and layer_class_to_replicate is not None:694 raise ValueError(695 "You can construct a Cache either from a list `layers` of all the predefined `CacheLayer`, or from a "696 "`layer_class_to_replicate`, in which case the Cache will append a new layer corresponding to "697 "`layer_class_to_replicate` for each new call to `update` with an idx not already in the Cache."698 )699 if layers is None and layer_class_to_replicate is None:700 raise ValueError(701 "You should provide exactly one of `layers` or `layer_class_to_replicate` to initialize a Cache."702 )703 self.layers = layers if layers is not None else []704 self.layer_class_to_replicate = layer_class_to_replicate705 self.offloading = offloading706 if self.offloading:707 self.only_non_sliding = offload_only_non_sliding708 self.prefetch_stream = torch.Stream() if _is_torch_greater_or_equal_than_2_7 else torch.cuda.Stream()709 710 def __repr__(self):711 return f"{self.__class__.__name__}(layers={self.layers})"712 713 def prefetch(self, layer_idx: int, only_non_sliding: bool = True):714 """715 Prefetch a given layer on its device. If `only_non_sliding` is True, it will try to prefetch only the layers716 which are non-sliding. If the `layer_idx` is outside the range, this will circle back to the first layers.717 Note that we use a non-default stream for this, to avoid blocking.718 """719 if only_non_sliding:720 # Try to find next non-sliding, starting at `layer_idx`721 try:722 layer_idx = layer_idx + self.is_sliding[layer_idx:].index(False)723 # In this case, we need to circle back to the beginning724 except ValueError:725 layer_idx = self.is_sliding.index(False)726 else:727 layer_idx = layer_idx if layer_idx < len(self.layers) else 0728 729 # Prefetch730 with self.prefetch_stream if _is_torch_greater_or_equal_than_2_7 else torch.cuda.stream(self.prefetch_stream):731 self.layers[layer_idx].prefetch()732 733 def offload(self, layer_idx: int, only_non_sliding: bool = True):734 """735 Offload a given `layer_idx`. If `only_non_sliding` is True, it will offload `layer_idx` only if it is a736 non-sliding layer. Note that we do it on the default stream, so that we ensure all earlier737 computation in the layer's `update` methods are finished.738 """739 if not (only_non_sliding and self.is_sliding[layer_idx]):740 self.layers[layer_idx].offload()741 742 def update(743 self,744 key_states: torch.Tensor,745 value_states: torch.Tensor,746 layer_idx: int,747 cache_kwargs: Optional[dict[str, Any]] = None,748 ) -> tuple[torch.Tensor, torch.Tensor]:749 """750 Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.751 752 Parameters:753 key_states (`torch.Tensor`):754 The new key states to cache.755 value_states (`torch.Tensor`):756 The new value states to cache.757 layer_idx (`int`):758 The index of the layer to cache the states for.759 cache_kwargs (`dict[str, Any]`, *optional*):760 Additional arguments for the cache subclass. These are specific to each subclass and allow new types of761 cache to be created.762 763 Return:764 A tuple containing the updated key and value states.765 """766 # In this case, the `layers` were not provided, and we must append as much as `layer_idx`767 if self.layer_class_to_replicate is not None:768 while len(self.layers) <= layer_idx:769 self.layers.append(self.layer_class_to_replicate())770 771 if self.offloading:772 # Wait for the stream to finish if needed, and start prefetching the next layer773 torch.cuda.default_stream(key_states.device).wait_stream(self.prefetch_stream)774 self.prefetch(layer_idx + 1, self.only_non_sliding)775 776 keys, values = self.layers[layer_idx].update(key_states, value_states, cache_kwargs)777 778 if self.offloading:779 self.offload(layer_idx, self.only_non_sliding)780 781 return keys, values782 783 def early_initialization(784 self, batch_size: int, num_heads: int, head_dim: int, dtype: torch.dtype, device: torch.device785 ):786 """787 Initialize all the layers in advance (it's otherwise lazily initialized on the first `update` call).788 This is useful for our `export` recipes, as `export` needs everything in advance.789 """790 # Note that the initialization needs all dimensions (except -2), as well as device and dtype, so we use791 # this fake tensor approach. It has size 0 on the -2 dimension, so it does not allocate any data (it only792 # creates an empty tensor with correct shape, dtype and device), which is very efficient and practical793 fake_keys_tensor = torch.zeros((batch_size, num_heads, 0, head_dim), dtype=dtype, device=device)794 # Init all layers795 for layer in self.layers:796 layer.lazy_initialization(fake_keys_tensor)797 798 def get_seq_length(self, layer_idx: int = 0) -> int:799 """Returns the sequence length of the cache for the given layer."""800 if layer_idx >= len(self.layers):801 return 0802 return self.layers[layer_idx].get_seq_length()803 804 def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]:805 """806 Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for807 the given layer at `layer_idx`.808 The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer.809 """810 # For DynamicCache, where the layers are created at runtime -> if it was not yet created, the size is811 # simply the shape of `cache_position`812 if layer_idx >= len(self.layers):813 return cache_position.shape[0], 0814 return self.layers[layer_idx].get_mask_sizes(cache_position)815 816 def get_max_cache_shape(self, layer_idx: int = 0) -> int:817 """Returns maximum sequence length of the cache object. Dynamic caches do not have a maximum length."""818 # For DynamicCache, where the layers are created at runtime -> if it was not yet created, return -1819 # as DynamicLayer does820 if layer_idx >= len(self.layers):821 return -1822 return self.layers[layer_idx].get_max_cache_shape()823 824 def reset(self):825 """Recursively reset all layers tensors"""826 for layer_idx in range(len(self.layers)):827 self.layers[layer_idx].reset()828 829 def reorder_cache(self, beam_idx: torch.LongTensor):830 """Reorder the cache for beam search"""831 for layer_idx in range(len(self.layers)):832 self.layers[layer_idx].reorder_cache(beam_idx)833 834 def crop(self, max_length: int):835 """Crop the cache to the given length"""836 for layer_idx in range(len(self.layers)):837 self.layers[layer_idx].crop(max_length)838 839 def batch_repeat_interleave(self, repeats: int):840 """Repeat and interleave the cache"""841 for layer_idx in range(len(self.layers)):842 self.layers[layer_idx].batch_repeat_interleave(repeats)843 844 def batch_select_indices(self, indices: torch.Tensor):845 """Select indices from the cache"""846 for layer_idx in range(len(self.layers)):847 self.layers[layer_idx].batch_select_indices(indices)848 849 @property850 def max_batch_size(self) -> int:851 """Return the maximum batch size of the cache"""852 values = [layer.max_batch_size for layer in self.layers]853 if len(set(values)) > 1:854 raise ValueError(f"Max batch size is not consistent across layers: {values}")855 return values[0]856 857 @property858 def max_cache_len(self) -> int:859 """Return the maximum cache length of the cache"""860 values = [layer.max_cache_len for layer in self.layers]861 return max(values)862 863 @property864 def is_compileable(self) -> bool:865 """Return whether the cache is compileable"""866 # For DynamicCache dispatching the layers lazily (otherwise, all([]) is True)867 if len(self.layers) == 0:868 return False869 return all(layer.is_compileable for layer in self.layers)870 871 @property872 def is_initialized(self) -> bool:873 """Return whether the cache data is initialized"""874 return len(self.layers) > 0 and all(layer.is_initialized for layer in self.layers)875 876 @property877 def is_sliding(self) -> list[bool]:878 """Return whether the layers of the cache are sliding window"""879 return [getattr(layer, "is_sliding", False) for layer in self.layers]880 881 def __getitem__(self, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor]:882 """883 Support for backwards-compatible `past_key_values` indexing, e.g. `past_key_values[0][0].shape[2]` to get the884 sequence length.885 """886 if layer_idx < len(self.layers):887 return self.layers[layer_idx].keys, self.layers[layer_idx].values888 else:889 raise KeyError(890 f"Cache only has {len(self.layers)} layers, attempted to access layer with index {layer_idx}"891 )892 893 def __iter__(self):894 """895 Support for backwards-compatible `past_key_values` iteration, e.g. `for x in past_key_values:` to iterate over896 keys and values897 """898 for layer_idx in range(len(self)):899 yield (self.layers[layer_idx].keys, self.layers[layer_idx].values)900 901 def __len__(self):902 """903 This value corresponds to the number of layers in the model.904 """905 # Note: for DynamicCache, layers are initialized lazily, so this will not be accurate before the first906 # forward through all the layers907 return len(self.layers)908 909 910class DynamicCache(Cache):911 """912 A cache that grows dynamically as more tokens are generated. This is the default for generative models.913 It stores the key and value states as a list of `CacheLayer`, one for each layer. The expected shape for each tensor914 in the `CacheLayer`s is `[batch_size, num_heads, seq_len, head_dim]`.915 If a config is passed, it will additionally check for sliding or hybrid cache structure, greatly reducing the916 memory requirement of the cached tensors to `[batch_size, num_heads, min(seq_len, sliding_window), head_dim]`.917 918 See `Cache` for details on common methods that are implemented by all cache classes.919 920 Args:921 ddp_cache_data (`Iterable[tuple[torch.Tensor, torch.Tensor]]`, *optional*):922 It was originally added for compatibility with `torch.distributed` (DDP). In a nutshell, it is923 `map(gather_map, zip(*caches))`, i.e. each item in the iterable contains the key and value states924 for a layer gathered across replicas by torch.distributed (shape=[global batch size, num_heads, seq_len, head_dim]).925 Note: it needs to be the 1st arg as well to work correctly926 config (`PretrainedConfig`, *optional*):927 The config of the model for which this Cache will be used. If passed, it will be used to check for sliding928 or hybrid layer structure, greatly reducing the memory requirement of the cached tensors to929 `[batch_size, num_heads, min(seq_len, sliding_window), head_dim]`.930 offloading (`bool`, *optional*, defaults to `False`):931 Whether to perform offloading of the layers to `cpu`, to save GPU memory.932 offload_only_non_sliding (`bool`, *optional*, defaults to `False`):933 If `offloading` is `True`, this further decides if only the non-sliding layers will be offloaded (because934 usually the sliding layers are small in size, so there is no need to offload them, and skipping it is faster).935 936 Example:937 938 ```python939 >>> from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache940 941 >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")942 >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")943 944 >>> inputs = tokenizer(text="My name is Qwen2", return_tensors="pt")945 946 >>> # Prepare a cache class and pass it to model's forward947 >>> past_key_values = DynamicCache(config=model.config)948 >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)949 >>> outputs.past_key_values # access cache filled with key/values from generation950 ```951 """952 953 def __init__(954 self,955 ddp_cache_data: Optional[Iterable[tuple[torch.Tensor, torch.Tensor]]] = None,956 config: Optional[PretrainedConfig] = None,957 offloading: bool = False,958 offload_only_non_sliding: bool = False,959 ):960 layers = []961 # If a config is passed, use it to infer the layer types and initialize accordingly962 if config is not None:963 decoder_config = config.get_text_config(decoder=True)964 sliding_window = getattr(decoder_config, "sliding_window", None) or getattr(965 decoder_config, "attention_chunk_size", None966 )967 layer_types = getattr(decoder_config, "layer_types", None)968 if layer_types is None:969 layer_types = [970 "sliding_attention" if sliding_window is not None else "full_attention"971 for _ in range(decoder_config.num_hidden_layers)972 ]973 # Some models have shared layers thus no cache is needed for them (e.g. Gemma3n)974 if hasattr(decoder_config, "num_kv_shared_layers"):975 layer_types = layer_types[: -decoder_config.num_kv_shared_layers]976 977 for layer_type in layer_types:978 # From a cache point of view, both sliding and chunked are the same in how they should behave and how many979 # states they should return - only the mask changes to make them different at the end!980 if layer_type in ("sliding_attention", "chunked_attention"):981 layers.append(DynamicSlidingWindowLayer(sliding_window=sliding_window))982 else:983 layers.append(DynamicLayer())984 985 # In this case, use the passed data to already fill in the Cache986 if ddp_cache_data is not None:987 # Init all the layers with the data988 for layer_idx, (key_states, value_states) in enumerate(ddp_cache_data):989 # If the config was not passed above, initialize a DynamicLayer for each entry of the ddp_data990 if config is None:991 layers.append(DynamicLayer())992 # Update the layer with the data993 _, _ = layers[layer_idx].update(key_states, value_states)994 995 # If neither of config nor ddp_data was passed, then simply lazy init a full cache of DynamicLayer996 if len(layers) == 0:997 super().__init__(998 layer_class_to_replicate=DynamicLayer,999 offloading=offloading,1000 offload_only_non_sliding=offload_only_non_sliding,1001 )1002 else:1003 super().__init__(layers=layers, offloading=offloading, offload_only_non_sliding=offload_only_non_sliding)1004 1005 def to_legacy_cache(self) -> tuple[tuple[torch.Tensor, torch.Tensor]]:1006 """1007 Converts the `Cache` instance into the its equivalent in the legacy cache format. Used for1008 backward compatibility.1009 """1010 legacy_cache = ()1011 for layer in self.layers:1012 legacy_cache += ((layer.keys, layer.values),)1013 return legacy_cache1014 1015 @classmethod1016 def from_legacy_cache(cls, past_key_values: tuple[tuple[torch.Tensor, torch.Tensor]]) -> "DynamicCache":1017 """1018 Converts a cache in the legacy cache format into an equivalent `Cache`. Used for1019 backward compatibility.1020 """1021 cache = cls()1022 if past_key_values is None:1023 logger.warning_once("past_key_values should not be None in from_legacy_cache()")1024 if past_key_values is not None:1025 for layer_idx in range(len(past_key_values)):1026 key_states, value_states = past_key_values[layer_idx]1027 cache.update(key_states, value_states, layer_idx)1028 return cache1029 1030 1031class StaticCache(Cache):1032 """1033 Static Cache class to be used with `torch.compile(model)` and `torch.export()`. It will check the `config`1034 for potential hybrid cache structure, and initialize each layer accordingly.1035 1036 See `Cache` for details on common methods that are implemented by all cache classes.1037 1038 Args:1039 config (`PretrainedConfig`):1040 The config of the model for which this Cache will be used. It will be used to check for sliding1041 or hybrid layer structure, and initialize each layer accordingly.1042 max_cache_len (`int`):1043 The maximum number of tokens that this Cache should hold.1044 offloading (`bool`, *optional*, defaults to `False`):1045 Whether to perform offloading of the layers to `cpu`, to save GPU memory.1046 offload_only_non_sliding (`bool`, *optional*, defaults to `True`):1047 If `offloading` is `True`, this further decides if only the non-sliding layers will be offloaded (because1048 usually the sliding layers are small in size, so there is no need to offload them, and skipping it is faster).1049 1050 Example:1051 1052 ```python1053 >>> from transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache1054 1055 >>> model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf")1056 >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")1057 1058 >>> inputs = tokenizer(text="My name is Llama", return_tensors="pt")1059 1060 >>> # Prepare a cache class and pass it to model's forward1061 >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate1062 >>> max_generated_length = inputs.input_ids.shape[1] + 101063 >>> past_key_values = StaticCache(config=model.config, max_cache_len=max_generated_length)1064 >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)1065 >>> outputs.past_key_values # access cache filled with key/values from generation1066 StaticCache()1067 ```1068 """1069 1070 # Pass-in kwargs as well to avoid crashing for BC (it used more arguments before)1071 def __init__(1072 self,1073 config: PretrainedConfig,1074 max_cache_len: int,1075 offloading: bool = False,1076 offload_only_non_sliding: bool = True,1077 **kwargs,1078 ):1079 config = config.get_text_config(decoder=True)1080 layer_types = getattr(config, "layer_types", None)1081 # If `layer_types` is not explicitly provided, infer if the model is fully sliding1082 if layer_types is None:1083 if getattr(config, "sliding_window", None) is not None:1084 layer_types = ["sliding_attention" for _ in range(config.num_hidden_layers)]1085 elif getattr(config, "attention_chunk_size", None) is not None:1086 layer_types = ["chunked_attention" for _ in range(config.num_hidden_layers)]1087 else:1088 layer_types = ["full_attention" for _ in range(config.num_hidden_layers)]1089 # Some models have shared layers thus no cache is needed for them (e.g. Gemma3n)1090 if hasattr(config, "num_kv_shared_layers"):1091 layer_types = layer_types[: -config.num_kv_shared_layers]1092 1093 layers = []1094 for layer_type in layer_types:1095 if layer_type == "sliding_attention":1096 layer = StaticSlidingWindowLayer(max_cache_len=max_cache_len, sliding_window=config.sliding_window)1097 elif layer_type == "chunked_attention":1098 # From a cache point of view, both sliding and chunked are the same in how they should behave and how many1099 # states they should return - only the mask changes to make them different at the end!1100 layer = StaticSlidingWindowLayer(1101 max_cache_len=max_cache_len, sliding_window=config.attention_chunk_size1102 )1103 else:1104 layer = StaticLayer(max_cache_len=max_cache_len)1105 layers.append(layer)1106 1107 super().__init__(layers=layers, offloading=offloading, offload_only_non_sliding=offload_only_non_sliding)1108 1109 1110class QuantizedCache(Cache):1111 """1112 A quantizer cache similar to what is described in the1113 [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache paper](https://huggingface.co/papers/2402.02750).1114 It allows the model to generate longer sequence length without allocating too much memory for keys and values1115 by applying quantization.1116 The cache has two types of storage, one for original precision and one for the1117 quantized cache. A `residual length` is set as a maximum capacity for the original precision cache. When the1118 length goes beyond maximum capacity, the original precision cache is discarded and moved into the quantized cache.1119 The quantization is done per-channel with a set `q_group_size` for both keys and values, in contrast to what was1120 described in the paper.1121 1122 See `Cache` for details on common methods that are implemented by all cache classes.1123 1124 Args:1125 backend (`str`):1126 The quantization backend to use. One of `("quanto", "hqq").1127 config (`PretrainedConfig`):1128 The config of the model for which this Cache will be used.1129 nbits (`int`, *optional*, defaults to 4):1130 The number of bits for quantization.1131 axis_key (`int`, *optional*, defaults to 0):1132 The axis on which to quantize the keys.1133 axis_value (`int`, *optional*, defaults to 0):1134 The axis on which to quantize the values.1135 q_group_size (`int`, *optional*, defaults to 64):1136 Quantization is done per-channel according to a set `q_group_size` for both keys and values.1137 residual_length (`int`, *optional*, defaults to 128):1138 Maximum capacity for the original precision cache1139 """1140 1141 def __init__(1142 self,1143 backend: str,1144 config: PretrainedConfig,1145 nbits: int = 4,1146 axis_key: int = 0,1147 axis_value: int = 0,1148 q_group_size: int = 64,1149 residual_length: int = 128,1150 ):1151 if backend == "quanto":1152 layer_class = QuantoQuantizedLayer1153 elif backend == "hqq":1154 layer_class = HQQQuantizedLayer1155 else:1156 raise ValueError(f"Unknown quantization backend `{backend}`")1157 1158 config = config.get_text_config(decoder=True)1159 layers = [1160 layer_class(nbits, axis_key, axis_value, q_group_size, residual_length)1161 for _ in range(config.num_hidden_layers)1162 ]1163 super().__init__(layers=layers)1164 1165 1166class EncoderDecoderCache(Cache):1167 """1168 Base, abstract class for all encoder-decoder caches. Can be used to hold combinations of self-attention and1169 cross-attention caches.1170 1171 See `Cache` for details on common methods that are implemented by all cache classes.1172 1173 Args:1174 caches (`Iterable`):1175 Usually an iterable of length 2, containing 2 `Cache` objects, the first one for self-attention, the1176 second one for cross-attention. Can optionally also be an iterable of length 1, containing a1177 `tuple[tuple[torch.Tensor]]` (usually used for compatibility with torch dp and ddp).1178 1179 Example:1180 1181 ```python1182 >>> from transformers import AutoProcessor, AutoModelForCausalLM, DynamicCache, EncoderDecoderCache1183 1184 >>> model = AutoModelForCausalLM.from_pretrained("openai/whisper-small")1185 >>> processor = AutoProcessor.from_pretrained("openai/whisper-small")1186 1187 >>> inputs = processor(audio=YOUR-AUDIO, return_tensors="pt")1188 1189 >>> # Prepare cache classes for encoder and decoder and pass it to model's forward1190 >>> self_attention_cache = DynamicCache(config=self.config)1191 >>> cross_attention_cache = DynamicCache(config=self.config)1192 >>> past_key_values = EncoderDecoderCache(self_attention_cache, cross_attention_cache)1193 >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)1194 >>> outputs.past_key_values # access cache filled with key/values from generation1195 EncoderDecoderCache()1196 ```1197 """1198 1199 def __init__(self, *caches) -> None:1200 # For dp and ddp support, if only one argument is passed, it should be an iterable of tuples of tensors