| import torch |
| from typing import Optional, Tuple, Dict, Any, List |
|
|
|
|
| class HybridMambaAttentionDynamicCache: |
| def __init__(self, config, batch_size, dtype=torch.bfloat16, device=None): |
| self.dtype = dtype |
| self.layers_block_type = config.layers_block_type |
| self.device = device |
| self.has_previous_state = False |
| self.max_batch_size = batch_size |
| self.sequence_len_offset = 0 |
| self.batch_size_offset = 0 |
|
|
| self.transformer_layers = [] |
| for i in range(len(self.layers_block_type)): |
| name = self.layers_block_type[i] |
| layer_type = name.split("_")[0] if "_" in name else name |
| if layer_type in ("Transformer", "HelixSeekMLA"): |
| self.transformer_layers.append(i) |
|
|
| num_layers = len(config.layers_block_type) |
| self.conv_states = [None for _ in range(num_layers)] |
| self.recurrent_states = [None for _ in range(num_layers)] |
| self.key_cache = [None for _ in range(num_layers)] |
| self.value_cache = [None for _ in range(num_layers)] |
|
|
| def update( |
| self, |
| key_states: torch.Tensor, |
| value_states: torch.Tensor, |
| layer_idx: int, |
| cache_kwargs: Optional[Dict[str, Any]] = None, |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| if self.key_cache[layer_idx] is None: |
| self.key_cache[layer_idx] = key_states |
| self.value_cache[layer_idx] = value_states |
| else: |
| self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2) |
| self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2) |
|
|
| return self.key_cache[layer_idx], self.value_cache[layer_idx] |
|
|
| def reorder_cache(self, beam_idx: torch.LongTensor): |
| for layer_idx in range(len(self.key_cache)): |
| if self.key_cache[layer_idx] is not None: |
| device = self.key_cache[layer_idx].device |
| beam_idx_device = beam_idx.to(device) |
| self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx_device) |
| self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx_device) |
|
|
| if self.conv_states[layer_idx] is not None: |
| device = self.conv_states[layer_idx][0].device |
| beam_idx_device = beam_idx.to(device) |
| q_conv, k_conv, v_conv = self.conv_states[layer_idx] |
| self.conv_states[layer_idx] = ( |
| q_conv.index_select(0, beam_idx_device), |
| k_conv.index_select(0, beam_idx_device), |
| v_conv.index_select(0, beam_idx_device), |
| ) |
|
|
| if self.recurrent_states[layer_idx] is not None: |
| device = self.recurrent_states[layer_idx].device |
| beam_idx_device = beam_idx.to(device) |
| self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].index_select(0, beam_idx_device) |
|
|
| def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: |
| if self.transformer_layers: |
| if layer_idx is None or layer_idx not in self.transformer_layers: |
| layer_idx = self.transformer_layers[0] |
| if layer_idx >= len(self.key_cache) or self.key_cache[layer_idx] is None: |
| return 0 |
| return self.key_cache[layer_idx].shape[-2] |
| for states in (self.conv_states, self.recurrent_states): |
| for s in states: |
| if s is not None: |
| if isinstance(s, tuple): |
| s = s[0] |
| return s.shape[1] if s.dim() >= 2 else s.shape[0] |
| return 0 |
|
|
| def get_max_length(self) -> Optional[int]: |
| return None |
|
|
| def __len__(self) -> int: |
| return len(self.layers_block_type) |
|
|
| def to_legacy_cache(self): |
| raise NotImplementedError("HybridMambaAttentionDynamicCache does not have a legacy cache equivalent.") |
|
|
| @classmethod |
| def from_legacy_cache(cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None): |
| raise NotImplementedError("HybridMambaAttentionDynamicCache does not have a legacy cache equivalent.") |
|
|
|
|
| class HelixSeekDeltaDynamicCache: |
| is_compileable = False |
|
|
| def __init__(self, config, batch_size=1, dtype=torch.bfloat16, device=None): |
| self.config = config |
| self.dtype = dtype |
| self.device = device |
| self.batch_size = batch_size |
| self._has_previous_state = False |
|
|
| num_layers = len(config.layers_block_type) |
| layer_types = [] |
| for i in range(num_layers): |
| layer_name = config.layers_block_type[i] if hasattr(config, 'layers_block_type') else "" |
| if layer_name.startswith("HelixSeekDelta"): |
| layer_types.append("linear_attention") |
| else: |
| layer_types.append("full_attention") |
| self.layer_types = layer_types |
|
|
| self.transformer_layers = [ |
| i for i in range(num_layers) if self.layer_types[i] == "full_attention" |
| ] |
|
|
| linear_layers = [i for i in range(num_layers) if self.layer_types[i] == "linear_attention"] |
| self.last_linear_layer = linear_layers[-1] if linear_layers else -1 |
|
|
| self.conv_states = [None for _ in range(num_layers)] |
| self.recurrent_states = [None for _ in range(num_layers)] |
| self.key_cache = [None for _ in range(num_layers)] |
| self.value_cache = [None for _ in range(num_layers)] |
|
|
| def __len__(self): |
| return len(self.layer_types) |
|
|
| def update( |
| self, |
| key_states: torch.Tensor, |
| value_states: torch.Tensor, |
| layer_idx: int, |
| cache_kwargs: Optional[Dict[str, Any]] = None, |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| self.key_cache[layer_idx] = key_states if self.key_cache[layer_idx] is None else torch.cat([self.key_cache[layer_idx], key_states], dim=2) |
| self.value_cache[layer_idx] = value_states if self.value_cache[layer_idx] is None else torch.cat([self.value_cache[layer_idx], value_states], dim=2) |
| return self.key_cache[layer_idx], self.value_cache[layer_idx] |
|
|
| def reorder_cache(self, beam_idx: torch.LongTensor): |
| for layer_idx in range(len(self.key_cache)): |
| if self.key_cache[layer_idx] is not None: |
| device = self.key_cache[layer_idx].device |
| beam_idx_device = beam_idx.to(device) |
| self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx_device) |
| self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx_device) |
|
|
| if self.conv_states[layer_idx] is not None: |
| device = self.conv_states[layer_idx][0].device |
| beam_idx_device = beam_idx.to(device) |
| q_conv, k_conv, v_conv = self.conv_states[layer_idx] |
| self.conv_states[layer_idx] = ( |
| q_conv.index_select(0, beam_idx_device), |
| k_conv.index_select(0, beam_idx_device), |
| v_conv.index_select(0, beam_idx_device), |
| ) |
|
|
| if self.recurrent_states[layer_idx] is not None: |
| device = self.recurrent_states[layer_idx].device |
| beam_idx_device = beam_idx.to(device) |
| self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].index_select(0, beam_idx_device) |
|
|
| def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: |
| if layer_idx is None: |
| layer_idx = 0 |
| if not self.transformer_layers: |
| layer_idx = 0 |
| elif layer_idx not in self.transformer_layers: |
| layer_idx = self.transformer_layers[0] |
| if layer_idx is None: |
| return 0 |
| if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None: |
| return 0 |
| return self.key_cache[layer_idx].shape[-2] |
|
|
| def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> Tuple[int, int]: |
| kv_offset = 0 |
| query_length = cache_position.shape[0] |
| past_seen_tokens = self.get_seq_length(layer_idx) |
| kv_length = query_length + past_seen_tokens |
| return kv_length, kv_offset |
|
|
| @property |
| def has_previous_state(self): |
| if self.last_linear_layer == -1: |
| return False |
| return self.conv_states[self.last_linear_layer] is not None |
|
|
