Module `Attention.Cache`

Key-value caches.

type 'b t = {

keys : (float, 'b) Nx.t;
values : (float, 'b) Nx.t;

}

The type for key-value caches with float dtype layout 'b: the projected keys and values of the positions seen so far, each of shape [| batch; num_heads; len; head_dim |]. Slots at positions not yet seen hold zeros and are never attended to.

val make : 
  ?batch:int ->
  num_heads:int ->
  head_dim:int ->
  len:int ->
  (float, 'b) Nx.dtype ->
  'b t

make ~num_heads ~head_dim ~len dtype is an empty cache of len slots: zero tensors of shape [| batch; num_heads; len; head_dim |]. batch defaults to 1. len bounds the total sequence length (prompt plus generated positions).

Raises Invalid_argument if any dimension is not positive.

val map : ('a 'c. ('a, 'c) Nx.t -> ('a, 'c) Nx.t) -> 'b t -> 'b t

map f c is c with f applied to c.keys and c.values, in that order. With map2 and iter it satisfies the Nx.Ptree.S contract at any fixed 'b, so caches can be leaves of a jitted step's parameter tree.

val map2 : 
  ('a 'c. ('a, 'c) Nx.t -> ('a, 'c) Nx.t -> ('a, 'c) Nx.t) ->
  'b t ->
  'b t ->
  'b t

map2 f c c' combines c and c' leafwise with f.

val iter : ('a 'c. ('a, 'c) Nx.t -> unit) -> 'b t -> unit

iter f c applies f to c.keys and c.values, in that order.

val astype : (float, 'c) Nx.dtype -> 'b t -> 'c t

astype dt c is c with c.keys and c.values cast to dt. Differentiable through Rune: gradients flow back at each original leaf's dtype, so an astype of float32 parameters inside a loss function yields float32 gradients.