Attention

Paper: Attention is all you need

See also Transformers, or distraction

Self-attention

Muti-head Attention

RingAttention

Paged Attention

paper

Used in conjunction with Continuous batching

Reduce memory usage of attention mechanism by swapping kv-cache in and out of memory. A block manager is similar to those of virtual memory in OS.

Essentially, it’s a form of paging, such that attention can be stored in contiguous memory. Partitions the KV cache of each sequence into KV blocks.

Given:

• each block contains KV vectors for fixed number of tokens, denoted as block size $B$.
• Key block $K_j=(k_{(j-1)B+1},\dots ,k_{jB})$
• Value block $V_j=(v_{(j-1)B+1},\dots ,v_{jB})$

$$A_{ij}=\frac{\exp (q_i^T{K}_j/\sqrt{d})}{{\sum }_{t=1}^{i//B}\exp (q_i^T{K}_t/\sqrt{d})},o_i=\sum _{j=1}^{i//B}{V}_j{A}_{ij}^T$$

where $A_{ij}=(a_{i,(j-1)B+1},\dots a_{i,jB})$ is row vector of attention score on j-th KV block.