transformer_lens.utilities.cache module

Cache utilities for TransformerLens.

This module contains utility functions for working with caches, particularly key-value caches used in transformer models.

transformer_lens.utilities.cache.get_pos_offset(past_kv_cache: TransformerLensKeyValueCache | None, batch_size: int) → int

Get position offset for KV cache.

Parameters:

  • past_kv_cache – Optional KV cache

  • batch_size – Batch size

Returns:

Position offset