laguna-martini / heapr /hf_cache.py
nikgeo's picture
Publish Laguna Martini grouped-pruning model card and reproducibility artifacts
6f11713 verified
Raw
History Blame Contribute Delete
378 Bytes
"""Helpers for fixed-shape Hugging Face KV caches."""
from __future__ import annotations
from typing import Any
def make_static_cache(model: Any, *, max_cache_len: int):
"""Create a Transformers StaticCache for direct model.forward calls."""
from transformers.cache_utils import StaticCache
return StaticCache(config=model.config, max_cache_len=max_cache_len)