Coverage for transformer_lens/utilities/cache.py: 100%

1"""Cache utilities for TransformerLens. 

2 

3This module contains utility functions for working with caches, particularly 

4key-value caches used in transformer models. 

5""" 

6 

7from typing import Optional 

8 

9from transformer_lens.cache.key_value_cache import TransformerLensKeyValueCache 

10 

11 

12def get_pos_offset(past_kv_cache: Optional[TransformerLensKeyValueCache], batch_size: int) -> int: 

13 """Get position offset for KV cache. 

14 

15 Args: 

16 past_kv_cache: Optional KV cache 

17 batch_size: Batch size 

18 

19 Returns: 

20 Position offset 

21 """ 

22 if past_kv_cache is None: 

23 return 0 

24 return past_kv_cache.entries[0].past_keys.shape[1]