A Single Key-Value Cache Entry for an Attention Layer

Stores cached key and value projections from an attention layer.

Format

An R6::R6Class object.

Details

Key and value tensors have shape (batch, num_heads, seq_len, head_dim).

Public fields

key: Cached key projections. A torch_tensor or NULL.
value: Cached value projections. A torch_tensor or NULL.

Methods

`KVCacheEntry$new()`

Create a new KVCacheEntry.

Usage

KVCacheEntry$new(key = NULL, value = NULL)

Arguments

key: Optional torch_tensor.
value: Optional torch_tensor.

`KVCacheEntry$is_valid()`

Check if this cache entry contains valid data.

Usage

KVCacheEntry$is_valid()

Returns

Logical scalar.

`KVCacheEntry$slice()`

Slice key/value along batch dimensions.

Returns a new KVCacheEntry with sliced tensors, or an empty entry if this entry is not valid.

Usage

KVCacheEntry$slice(indices, drop = FALSE)

Arguments

indices: Indices for slicing. Can be an integer vector, a torch_tensor, or use .. for all elements.
drop: drop unitary dimension

Returns

A new KVCacheEntry instance.

`KVCacheEntry$write()`

Write a batch slice into this entry.

Usage

KVCacheEntry$write(indices, other)

Arguments

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

Returns

Invisible self (for method chaining).

`KVCacheEntry$to()`

Move this entry to the given device and optionally cast dtype.

Returns a new KVCacheEntry.

Usage

KVCacheEntry$to(device, dtype = NULL)

Arguments

device: Device string, e.g. "cpu" or "cuda".
dtype: Optional torch_dtype.

Returns

A new KVCacheEntry instance.

`KVCacheEntry$clone()`

The objects of this class are cloneable with this method.

Usage

KVCacheEntry$clone(deep = FALSE)

Arguments

deep: Whether to make a deep clone.

Examples

if (FALSE) { # \dontrun{
  entry <- KVCacheEntry$new(
    key = torch_randn(2, 4, 10, 64),
    value = torch_randn(2, 4, 10, 64)
  )
  entry$is_valid()
} # }