Class KVState

axlearn/common/kv_cache/base_kv_cache.py:15–28

Represents key/value projections. Fields: k_proj: [batch, source_length, num_kv_heads, per_head_dim], Projected key tensor. v_proj: [batch, source_length, num_kv_heads, per_head_dim], Projected value tensor. key_positions: [batch, source_length], Positions of the keys in

class KVState(NamedTuple):
    """Represents key/value projections.

    Fields:
        k_proj: [batch, source_length, num_kv_heads, per_head_dim], Projected key tensor.
        v_proj: [batch, source_length, num_kv_heads, per_head_dim], Projected value tensor.
        key_positions: [batch, source_length], Positions of the keys in the batch.
        page_indices: [batch, max_pages_per_request], optional page indices for batched requests.
    """

    k_proj: Tensor
    v_proj: Tensor
    key_positions: Tensor
    page_indices: Optional[Tensor] = None


class BaseKVCache(BaseLayer):