skhavin/proactive-cache

1"""2utils.py — DynamicCache compatibility helpers.3 4Handles the transformers DynamicCache ↔ legacy tuple conversion cleanly5across transformers versions 4.38+.6"""7 8from __future__ import annotations9import torch10from typing import Tuple, Union11 12 13# Type aliases14KVTuple = Tuple[Tuple[torch.Tensor, torch.Tensor], ...]15 16 17def to_tuple_kv(past_key_values) -> KVTuple:18    """Normalize a DynamicCache or legacy tuple to a tuple of (k, v) pairs."""19    if hasattr(past_key_values, "to_legacy_cache"):20        return past_key_values.to_legacy_cache()21    return tuple(past_key_values)22 23 24def to_dynamic_cache(kv_tuple: KVTuple):25    """Convert a (k, v) tuple back to DynamicCache for models that require it."""26    try:27        from transformers import DynamicCache28        return DynamicCache.from_legacy_cache(kv_tuple)29    except (ImportError, AttributeError):30        # Older transformers — raw tuple is fine31        return kv_tuple32 33 34def get_device(model) -> torch.device:35    """Get the primary device of a model."""36    return next(model.parameters()).device37 38 39def get_num_layers(past_key_values) -> int:40    """Return the number of transformer layers in a KV cache."""41    kv = to_tuple_kv(past_key_values)42    return len(kv)43 44 45def get_seq_len(past_key_values) -> int:46    """Return the current sequence length stored in a KV cache."""47    kv = to_tuple_kv(past_key_values)48    if len(kv) == 0:49        return 050    # Shape: (batch, num_heads, seq_len, head_dim)51    return kv[0][0].shape[2]52