| """Shared Safin-1 cache utils for dense and MoE remote code.""" |
|
|
| from __future__ import annotations |
|
|
|
|
| from types import MethodType |
| import torch |
| from torch import nn |
| from transformers.cache_utils import DynamicCache |
| from transformers.models.qwen3_5_moe import modeling_qwen3_5_moe as qwen35_moe_modeling |
|
|
|
|
| def _new_qwen35_cache(config): |
| cache_cls = getattr(qwen35_moe_modeling, "Qwen3_5MoeDynamicCache", DynamicCache) |
| try: |
| cache = cache_cls(config=config) |
| except TypeError: |
| cache = cache_cls() |
| return _patch_hsd_cache_get_seq_length(cache) |
|
|
|
|
| def _patch_hsd_cache_get_seq_length(cache): |
| if cache is None: |
| return cache |
| _patch_hsd_cache_batch_ops(cache) |
| if getattr(cache, "_historical_hsd_get_seq_length_patched", False): |
| return cache |
| original_get_seq_length = getattr(cache, "get_seq_length", None) |
| if not callable(original_get_seq_length): |
| return cache |
|
|
| def get_seq_length(self, *args, **kwargs): |
| try: |
| return original_get_seq_length(*args, **kwargs) |
| except ValueError as exc: |
| message = str(exc) |
| if "get_seq_length" not in message or "LinearAttention" not in message: |
| raise |
| value = getattr(self, "hsd_seen_expanded_tokens", None) |
| if value is not None: |
| return _length_value_to_int(value) |
| value = getattr(self, "hsd_seen_text_tokens", None) |
| if value is not None: |
| return _length_value_to_int(value) |
| return 0 |
|
|
| cache.get_seq_length = MethodType(get_seq_length, cache) |
| cache._historical_hsd_get_seq_length_patched = True |
| return cache |
|
|
|
|
| def _cache_has_previous_state(past_key_values) -> bool: |
| if past_key_values is None: |
| return False |
| if _length_value_to_int(getattr(past_key_values, "hsd_seen_expanded_tokens", 0)) > 0: |
| return True |
| if _length_value_to_int(getattr(past_key_values, "hsd_seen_text_tokens", 0)) > 0: |
| return True |
| if _cache_text_length(past_key_values) > 0: |
| return True |
| has_previous_state = getattr(past_key_values, "has_previous_state", None) |
| if callable(has_previous_state): |
| try: |
| return bool(has_previous_state()) |
| except TypeError: |
| return bool(has_previous_state(0)) |
| if isinstance(has_previous_state, bool): |
| return has_previous_state |
| try: |
| return int(past_key_values.get_seq_length()) > 0 |
| except Exception: |
| return False |
|
|
|
|
| def _cache_text_length(past_key_values) -> int: |
| if past_key_values is None: |
| return 0 |
| value = getattr(past_key_values, "hsd_seen_text_tokens", None) |
| if value is not None: |
| return _length_value_to_int(value) |
| layers = getattr(past_key_values, "layers", None) |
| if layers is not None: |
| for layer in layers: |
| value = getattr(layer, "hsd_seen_text_tokens", None) |
| if value is not None: |
| return _length_value_to_int(value) |
| layers = getattr(past_key_values, "_hsd_layer_states", None) |
| if layers is not None: |
| for layer in layers: |
| value = getattr(layer, "hsd_seen_text_tokens", None) |
| if value is not None: |
| return _length_value_to_int(value) |
| try: |
| return int(past_key_values.get_seq_length()) |
| except Exception: |
| return 0 |
|
|
|
|
| def _cache_expanded_length(past_key_values) -> int: |
| if past_key_values is None: |
| return 0 |
| value = getattr(past_key_values, "hsd_seen_expanded_tokens", None) |
| if value is not None: |
| return _length_value_to_int(value) |
| try: |
| value = int(past_key_values.get_seq_length()) |
| if value > 0: |
| return value |
| except Exception: |
| pass |
| text_len = _cache_text_length(past_key_values) |
| return text_len |
|
|
|
|
| def _length_value_to_int(value) -> int: |
| if value is None: |
| return 0 |
| if torch.is_tensor(value): |
| if value.numel() == 0: |
| return 0 |
| return int(value.detach().max().item()) |
| return int(value or 0) |
|
|
|
|
| def _cache_lengths( |
| past_key_values, |
| *, |
| base_attr: str, |
| per_row_attr: str, |
| batch_size: int, |
| device: torch.device, |
| ) -> torch.Tensor: |
| if past_key_values is None: |
| return torch.zeros(batch_size, dtype=torch.long, device=device) |
| for source in (past_key_values,): |
| value = getattr(source, per_row_attr, None) |
| if value is not None: |
| tensor = torch.as_tensor(value, dtype=torch.long, device=device).reshape(-1) |
| if tensor.numel() == batch_size: |
| return tensor |
| if tensor.numel() == 1: |
| return tensor.expand(batch_size) |
| for attr_name in ("layers", "_hsd_layer_states"): |
| layers = getattr(past_key_values, attr_name, None) |
| if layers is None: |
| continue |
| for layer in layers: |
| value = getattr(layer, per_row_attr, None) |
| if value is not None: |
| tensor = torch.as_tensor(value, dtype=torch.long, device=device).reshape(-1) |
| if tensor.numel() == batch_size: |
| return tensor |
| if tensor.numel() == 1: |
| return tensor.expand(batch_size) |
| value = getattr(layer, base_attr, None) |
| if value is not None: |
| scalar = _length_value_to_int(value) |
| return torch.full((batch_size,), scalar, dtype=torch.long, device=device) |
| value = getattr(past_key_values, base_attr, None) |
| if value is not None: |
| scalar = _length_value_to_int(value) |
| return torch.full((batch_size,), scalar, dtype=torch.long, device=device) |
| scalar = _cache_text_length(past_key_values) if base_attr == "hsd_seen_text_tokens" else _cache_expanded_length(past_key_values) |
| return torch.full((batch_size,), scalar, dtype=torch.long, device=device) |
|
|
|
|
| def _cache_text_lengths(past_key_values, batch_size: int, device: torch.device) -> torch.Tensor: |
| return _cache_lengths( |
| past_key_values, |
| base_attr="hsd_seen_text_tokens", |
| per_row_attr="hsd_seen_text_tokens_per_row", |
| batch_size=batch_size, |
| device=device, |
| ) |
|
|
|
|
| def _cache_expanded_lengths(past_key_values, batch_size: int, device: torch.device) -> torch.Tensor: |
| return _cache_lengths( |
| past_key_values, |
| base_attr="hsd_seen_expanded_tokens", |
| per_row_attr="hsd_seen_expanded_tokens_per_row", |
| batch_size=batch_size, |
| device=device, |
| ) |
|
|
|
|
| def _set_root_cache_lengths( |
| past_key_values, |
| *, |
| text_lengths: torch.Tensor, |
| expanded_lengths: torch.Tensor, |
| ) -> None: |
| if past_key_values is None: |
| return |
| text_cpu = text_lengths.detach().to(dtype=torch.long) |
| expanded_cpu = expanded_lengths.detach().to(dtype=torch.long) |
| past_key_values.hsd_seen_text_tokens_per_row = text_cpu |
| past_key_values.hsd_seen_expanded_tokens_per_row = expanded_cpu |
| past_key_values.hsd_seen_text_tokens = _length_value_to_int(text_cpu) |
| past_key_values.hsd_seen_expanded_tokens = _length_value_to_int(expanded_cpu) |
|
|
|
|
| class _RowFullAttentionCache: |
| def __init__(self, layer_idx: int, keys: torch.Tensor | None = None, values: torch.Tensor | None = None): |
| self.layer_idx = int(layer_idx) |
| self.keys = keys |
| self.values = values |
|
|
| def get_seq_length(self, layer_idx: int = 0) -> int: |
| |
| return int(self.keys.shape[-2]) if torch.is_tensor(self.keys) else 0 |
|
|
| def get_mask_sizes(self, query_length: int | torch.Tensor, layer_idx: int = 0) -> tuple[int, int]: |
| length = int(query_length.numel()) if torch.is_tensor(query_length) else int(query_length) |
| return self.get_seq_length() + length, 0 |
|
|
| def update(self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, *args, **kwargs): |
| if int(layer_idx) != self.layer_idx: |
| raise ValueError(f"row full-attention cache received layer {layer_idx}, expected {self.layer_idx}.") |
| if self.keys is None or self.values is None or self.keys.numel() == 0: |
| self.keys = key_states |
| self.values = value_states |
| else: |
| self.keys = torch.cat([self.keys.to(key_states.device), key_states], dim=-2) |
| self.values = torch.cat([self.values.to(value_states.device), value_states], dim=-2) |
| return self.keys, self.values |
|
|
|
|
| def _full_attention_layer_idx(layer: nn.Module) -> int: |
| attn = getattr(layer, "self_attn", None) |
| return int(getattr(attn, "layer_idx", getattr(layer, "_historical_hsd_layer_idx", getattr(layer, "layer_idx", 0)))) |
|
|
|
|
| def _slice_row_full_attention_cache(past_key_values, layer_idx: int, batch_idx: int, batch_size: int) -> _RowFullAttentionCache | None: |
| if past_key_values is None: |
| return None |
| layers = getattr(past_key_values, "layers", None) |
| if layers is None or len(layers) <= layer_idx: |
| return _RowFullAttentionCache(layer_idx) |
| layer_cache = layers[layer_idx] |
| keys = getattr(layer_cache, "keys", None) |
| values = getattr(layer_cache, "values", None) |
| if not torch.is_tensor(keys) or not torch.is_tensor(values) or keys.numel() == 0: |
| return _RowFullAttentionCache(layer_idx) |
| lengths = getattr(layer_cache, "hsd_text_cache_lengths_per_row", None) |
| if torch.is_tensor(lengths) and lengths.numel() > batch_idx: |
| row_len = int(lengths.reshape(-1)[batch_idx].item()) |
| else: |
| row_len = int(keys.shape[-2]) |
| return _RowFullAttentionCache( |
| layer_idx, |
| keys[batch_idx : batch_idx + 1, :, :row_len].contiguous(), |
| values[batch_idx : batch_idx + 1, :, :row_len].contiguous(), |
| ) |
|
|
|
|
| def _merge_row_full_attention_caches(past_key_values, layer_idx: int, row_caches: list[_RowFullAttentionCache | None]) -> None: |
| if past_key_values is None: |
| return |
| layers = getattr(past_key_values, "layers", None) |
| if layers is None or len(layers) <= layer_idx: |
| return |
| keyed = [cache for cache in row_caches if cache is not None and torch.is_tensor(cache.keys)] |
| if not keyed: |
| return |
| sample_keys = keyed[0].keys |
| sample_values = keyed[0].values |
| max_len = max(int(cache.keys.shape[-2]) if cache is not None and torch.is_tensor(cache.keys) else 0 for cache in row_caches) |
| batch_size = len(row_caches) |
| merged_keys = sample_keys.new_zeros(batch_size, sample_keys.shape[1], max_len, sample_keys.shape[-1]) |
| merged_values = sample_values.new_zeros(batch_size, sample_values.shape[1], max_len, sample_values.shape[-1]) |
| lengths = torch.zeros(batch_size, dtype=torch.long, device=merged_keys.device) |
| for batch_idx, cache in enumerate(row_caches): |
| if cache is None or not torch.is_tensor(cache.keys) or not torch.is_tensor(cache.values): |
| continue |
| row_len = int(cache.keys.shape[-2]) |
| merged_keys[batch_idx : batch_idx + 1, :, :row_len] = cache.keys.to(merged_keys.dtype) |
| merged_values[batch_idx : batch_idx + 1, :, :row_len] = cache.values.to(merged_values.dtype) |
| lengths[batch_idx] = row_len |
| layer_cache = layers[layer_idx] |
| layer_cache.keys = merged_keys |
| layer_cache.values = merged_values |
| layer_cache.is_initialized = True |
| layer_cache.dtype = merged_keys.dtype |
| layer_cache.device = merged_keys.device |
| layer_cache.hsd_text_cache_lengths_per_row = lengths |
|
|
|
|
| def _reorder_batch_tensor(value, beam_idx: torch.LongTensor, *, min_ndim=1): |
| """Select rows of a declared batch tensor, including expansion and pruning.""" |
| if torch.is_tensor(value) and value.ndim >= min_ndim: |
| return value.index_select(0, beam_idx.to(value.device)) |
| if isinstance(value, list): |
| return [_reorder_batch_tensor(item, beam_idx, min_ndim=min_ndim) for item in value] |
| if isinstance(value, tuple): |
| return tuple(_reorder_batch_tensor(item, beam_idx, min_ndim=min_ndim) for item in value) |
| return value |
|
|
|
|
| _PER_ROW_LENGTHS = ("hsd_seen_text_tokens_per_row", "hsd_seen_expanded_tokens_per_row", |
| "hsd_text_cache_lengths_per_row") |
| _ANCHOR_BATCH_FIELDS = ("hsd_anchor_states", "hsd_anchor_keys", "hsd_anchor_boundaries", |
| "hsd_anchor_valid_mask", "hsd_anchor_persistent_memory_mask") |
|
|
|
|
| def _hsd_cache_objects(cache): |
| seen = set() |
| for obj in [cache, *(getattr(cache, "_hsd_layer_states", None) or []), *(getattr(cache, "layers", None) or [])]: |
| if obj is not None and id(obj) not in seen: |
| seen.add(id(obj)) |
| yield obj |
|
|
|
|
| def _hsd_cache_batch_size(cache): |
| for obj in _hsd_cache_objects(cache): |
| for name in (*_PER_ROW_LENGTHS, "hsd_anchor_states", "conv_states", "recurrent_states", "keys"): |
| value = getattr(obj, name, None) |
| if torch.is_tensor(value) and value.ndim > 0: |
| return int(value.shape[0]) |
| for name in ("_hsd_conv_states", "_hsd_recurrent_states"): |
| for value in getattr(cache, name, []) or []: |
| if torch.is_tensor(value) and value.ndim > 0: |
| return int(value.shape[0]) |
| return None |
|
|
|
|
| def _validate_cache_indices(cache, indices): |
| if not torch.is_tensor(indices) or indices.ndim != 1 or indices.dtype not in (torch.int32, torch.int64): |
| raise ValueError("cache row indices must be a one-dimensional integer tensor") |
| batch_size = _hsd_cache_batch_size(cache) |
| if indices.numel() and (int(indices.min()) < 0 or |
| (batch_size is not None and int(indices.max()) >= batch_size)): |
| raise IndexError(f"cache row indices are out of range for batch size {batch_size}") |
|
|
|
|
| def _reorder_hsd_sidecar_cache(past_key_values, beam_idx: torch.LongTensor): |
| """Reorder MARCH sidecars only; upstream owns native KV/conv/recurrent layers.""" |
| if past_key_values is None: |
| return past_key_values |
| _validate_cache_indices(past_key_values, beam_idx) |
| for obj in _hsd_cache_objects(past_key_values): |
| for name in _PER_ROW_LENGTHS: |
| if hasattr(obj, name): |
| setattr(obj, name, _reorder_batch_tensor(getattr(obj, name), beam_idx)) |
| for name in _ANCHOR_BATCH_FIELDS: |
| if hasattr(obj, name): |
| |
| |
| setattr(obj, name, _reorder_batch_tensor(getattr(obj, name), beam_idx, min_ndim=2)) |
| for base in ("hsd_seen_text_tokens", "hsd_seen_expanded_tokens"): |
| lengths = getattr(obj, base + "_per_row", None) |
| if torch.is_tensor(lengths): |
| setattr(obj, base, _length_value_to_int(lengths)) |
| for name in ("_hsd_conv_states", "_hsd_recurrent_states"): |
| states = getattr(past_key_values, name, None) |
| if states is not None: |
| setattr(past_key_values, name, _reorder_batch_tensor(states, beam_idx)) |
| return past_key_values |
|
|
|
|
| def _patch_hsd_cache_batch_ops(cache): |
| """Give all HF cache row operations one native reorder and one sidecar reorder. |
| |
| LinearAttentionLayer lacks the upstream repeat/select methods. Implement |
| those through its native reorder_cache, which also handles size changes. |
| """ |
| if getattr(cache, "_march_batch_ops_patched", False): |
| return cache |
| original = getattr(cache, "reorder_cache", None) |
| if not callable(original): |
| return cache |
|
|
| def reorder_cache(self, indices): |
| _validate_cache_indices(self, indices) |
| |
| _reorder_hsd_sidecar_cache(self, indices) |
| return original(indices) |
|
|
| def batch_select_indices(self, indices): |
| return self.reorder_cache(indices) |
|
|
| def batch_repeat_interleave(self, repeats): |
| if not isinstance(repeats, int) or isinstance(repeats, bool) or repeats < 1: |
| raise ValueError("cache repeats must be a positive integer") |
| batch_size = _hsd_cache_batch_size(self) |
| if batch_size is None: |
| return None |
| indices = torch.arange(batch_size, dtype=torch.long).repeat_interleave(repeats) |
| return self.reorder_cache(indices) |
|
|
| cache.reorder_cache = MethodType(reorder_cache, cache) |
| cache.batch_select_indices = MethodType(batch_select_indices, cache) |
| cache.batch_repeat_interleave = MethodType(batch_repeat_interleave, cache) |
| cache._march_batch_ops_patched = True |
| return cache |
|
|
