Base Class for Key-Value Caches — KVCache • TabICL2

Provides common structure and operations for caches that store mappings from layer/block index to cached key-value projections.

Format

An R6::R6Class object.

Details

Layer indices are stored as character strings in the named list kv, since R lists cannot have integer names that start from 0.

Public fields

kv: Named list mapping layer index (character) to KVCacheEntry.

Methods

`KVCache$new()`

Create a new KVCache.

Usage

KVCache$new(kv = list())

Arguments

kv: Optional named list of KVCacheEntry instances.

`KVCache$is_populated()`

Check if this cache has valid entries.

Returns TRUE when the cache contains data (use_cache mode). Returns FALSE when the cache is empty (store_cache mode).

Usage

KVCache$is_populated()

Returns

Logical scalar.

`KVCache$slice()`

Slice all entries along batch dimensions.

Returns a new cache of the same class with sliced entries.

Usage

KVCache$slice(indices)

Arguments

indices: Indices for slicing.

Returns

A new KVCache instance.

`KVCache$write()`

Write batch-sliced entries into this pre-allocated cache.

Usage

KVCache$write(indices, other)

Arguments

indices: Indices for writing.
other: A KVCache instance to write from.

Returns

Invisible self (for method chaining).

`KVCache$to()`

Move all entries to the given device and optionally cast dtype.

Returns a new cache of the same class.

Usage

KVCache$to(device, dtype = NULL)

Arguments

device: Device string.
dtype: Optional torch_dtype.

Returns

A new KVCache instance.

`KVCache$preallocate()`

Pre-allocate entries in this cache based on shapes from a reference.

K/V tensors always have shape (*batch, num_heads, seq_len, head_dim). This method keeps the last three dimensions from the reference entry and prepends batch_shape as the leading dimensions.

Usage

KVCache$preallocate(reference, batch_shape, device = "cpu", dtype = NULL)

Arguments

reference: A KVCache from a single batch whose entry shapes are used as a template.
batch_shape: Integer vector. The full batch shape for leading dims.
device: Device on which to allocate tensors. Default "cpu".
dtype: Optional torch_dtype. If NULL, uses reference entry dtype.

Returns

Invisible self (for method chaining).

`KVCache$clone()`

The objects of this class are cloneable with this method.

Usage

KVCache$clone(deep = FALSE)

Arguments

deep: Whether to make a deep clone.

Examples

if (FALSE) { # \dontrun{
  cache <- KVCache$new()
  cache$kv[["1"]] <- KVCacheEntry$new(
    key = torch_randn(2, 4, 10, 64),
    value = torch_randn(2, 4, 10, 64)
  )
  cache$is_populated()
} # }