PagedAttention: from interface to kernal

1 Overview

PagedAttention灵感来自于操作系统中虚拟内存和分页的经典思想，它可以允许在非连续空间立存储连续的KV张量。具体来说，PagedAttention把每个序列的KV缓存进行了分块，每个块包含固定长度的token，而在计算attention时可以高效地找到并获取那些块。

在这里插入图片描述

2 Block management

相比于RaggedAttention，PagedAttention其实就是维护了一个逻辑block到物理block的映射。接下来，我们结合代码，看一下是vLLM如何实现PagedAttention中Block管理的。

2.1 LogicalTokenBlock and PhysicalTokenBlock

先看一下LogicalTokenBlock 和 PhysicalTokenBlock的数据结构，比较简单，看代码注释就好了。

class LogicalTokenBlock:
    """
    Logical blocks are used to represent the states of the corresponding physical blocks in the KV cache.
    """
    def __init__(
        self,
        block_number: int,
        block_size: int,
    ) -> None:
        self.block_number = block_number	// LogicalTokenBlock id, 由零递增
        self.block_size = block_size	// 最多容纳token数量（block大小），默认为 16
        self.token_ids = [_BLANK_TOKEN_ID] * block_size	// 初始化为TOKEN_ID为-1，表示无效token
        self.num_tokens = 0	// 此block中存储的有效token的数量
        
    def append_tokens(self, token_ids: List[int]) -> None:
        """核心api, LogicalTokenBlock中的token只能从左到右append"""
        assert len(token_ids) <= self.get_num_empty_slots()
        curr_idx = self.num_tokens
        self.token_ids[curr_idx:curr_idx + len(token_ids)] = token_ids
        self.num_tokens += len(token_ids)

class PhysicalTokenBlock:
    """Represents the state of a block in the KV cache."""
    def __init__(
        self,
        device: Device,
        block_number: int,
        block_size: int,
    ) -> None:
        self.device = device	// 记录PhysicalTokenBlock物理存储位置，GPU or CPU
        self.block_number = block_number	// PhysicalTokenBlock id, 由零递增
        self.block_size = block_size	// block大小，与LogicalTokenBlock中block_size相等，默认为 16
        self.ref_count = 0	// 引用计数，beam search 里面同一个block可能会被多个sequence引用

2.2 BlockAllocator

BlockAllocator是用来管理空闲block的数据结构。其中最核心的数据结构是一个链表free_blocks，存储空闲的PhysicalTokenBlock。

# Mapping: logical block number -> physical block.
BlockTable = List[PhysicalTokenBlock]

class BlockAllocator:
    """Manages free physical token blocks for a device.

    The allocator maintains a list of free blocks and allocates a block when
    requested. When a block is freed, its reference count is decremented. If
    the reference count becomes zero, the block is added back to the free list.
    """

    def __init__(
        self,
        device: Device,
        block_size: int,
        num_blocks: int,
    ) -> None:
        self.device = device	// 维护的blocks存储位置，GPU or CPU
        self.block_size = block_size	// block大小，默认为 16
        self.num_blocks = num_blocks	// 

        # Initialize the free blocks.
        self.free_blocks: BlockTable = []	// 一个链表，存储空闲的PhysicalTokenBlock
        for i in range(num_blocks):
            block = PhysicalTokenBlock(device=device,
                                       block_number=i,
                                       block_size=block_size)
            self.free_blocks.append(block)

2.3 BlockSpaceManager

BlockSpaceManager是维护BlockSpace的核心数据结构，维护着LogicalTokenBlock到PhysicalTokenBlock的映射。在vLLM初始化的时候，首先会获取系统可用的CPU端和GPU端存储容量，从而计算出可用的block数量。

gpu_allocator
- 将系统可用的GPU端存储按照block进行管理。
- 其中维护一个链表free_blocks，存储空闲的PhysicalTokenBlock。
cpu_allocator
- 将系统可用的CPU端存储按照block进行管理。
- 当GPU显存不够用时，会将存储在GPU显存的block换出至CPU。
watermark_blocks
- 其实就是设置一个阈值，当gpu_allocator内空闲block数量小于watermark_blocks，停止分配新的sequence。
- 以避免频繁的缓存逐出。
block_tables
- 存储sequence的BlockTable
- BlockTable 维护者 logical block number -> physical block 的映射
- 假设一个sequence使用了N个logical block，那么，logical block number 为0~N-1

class BlockSpaceManager:
    """Manages the mapping between logical and physical token blocks."""

    def __init__(
        self,
        block_size: int,
        num_gpu_blocks: int,
        num_cpu_blocks: int,
        watermark: float = 0.01
    ) -> None:
        self.block_size = block_size
        self.num_total_gpu_blocks = num_gpu_blocks
        self.num_total_cpu_blocks = num_cpu_blocks

        self.watermark = watermark
        assert watermark >= 0.0

        self.watermark_blocks = int(watermark * num_gpu_blocks)
        self.gpu_allocator = BlockAllocator(Device.GPU, block_size,
                                            num_gpu_blocks)
        self.cpu_allocator = BlockAllocator(Device.CPU, block_size,
                                            num_cpu_blocks)
        # Mapping: seq_id -> BlockTable.
        self.block_tables: Dict[int, BlockTable] = {}

2.4 Sequence

一个Sequence对应一个推理上下文。

logical_token_blocks
- 一个链表，维护着Sequence使用的LogicalTokenBlock
- 链表中LogicalTokenBlock的logical block number 依次为0~N-1
- 根据logical block number就可以在BlockSpaceManager中的block_tables获取physical block

class Sequence:
    """Stores the data, status, and block information of a sequence.

    Args:
        seq_id: The ID of the sequence.
        prompt: The prompt of the sequence.
        prompt_token_ids: The token IDs of the prompt.
        block_size: The block size of the sequence. Should be the same as the
            block size used by the block manager and cache engine.
    """

    def __init__(
        self,
        seq_id: int,
        prompt: str,
        prompt_token_ids: List[int],
        block_size: int,
    ) -> None:
        self.seq_id = seq_id
        self.prompt = prompt
        self.block_size = block_size
        self.logical_token_blocks: List[LogicalTokenBlock] = []
        # Initialize the logical token blocks with the prompt token ids.
        self._append_tokens_to_blocks(prompt_token_ids)

3 PagedAttention kernal

Transformer推理过程中最重要的算子是attention算子。当拿到物理block，就可以获取到当前context中token对应的kv vector，然后进行attention操作。

这是一张Multi-head attention（NUM_HEAD = 2）的示意图，q向量与kv向量进行attention操作，产生下一个token。

在这里插入图片描述

3.1 attention 伪代码

attention 的C语言串行代码非常简单，容易理解。

q 向量与 MAX_SEQ个 k向量分别进行向量内积计算。

计算出来MAX_SEQ个logit，代表了 q 向量和 k 向量之间的相关性。
对MAX_SEQ个logit进行softmax
q 向量和 k 向量之间的相关性作用于v向量，获取最终的ans向量。

typedef struct _query
{
    vector vectors[NUM_HEAD];
} query;
typedef query answer;
typedef struct _kv_cache
{
    uint32_t num_token;
    vector k_cache[NUM_HEAD][MAX_SEQ_LEN];
    vector v_cache[NUM_HEAD][MAX_SEQ_LEN];
} kv_cache;

void multi_head_attention(query *q, kv_cache *kv, answer *ans)
{
    int32_t logits[MAX_SEQ_LEN];
    vector *head_ans;

    answer_init(ans);
    // NUM_HEAD个attention操作相互独立
    for (int i = 0; i < NUM_HEAD; i++)
    {
        // q 向量与 MAX_SEQ个 k向量分别进行向量内积计算。logits[j]代表了q->vectors[i]和kv->k_cache[i][j]之间的相关性
        for (int j = 0; j < MAX_SEQ_LEN; j++)
        {
            logits[j] = vector_dot(&(q->vectors[i]), &(kv->k_cache[i][j]));
        }
        // 对MAX_SEQ个logit进行softmax
        soft_max(logits);
        head_ans = &(ans->vectors[i]);
        // logits[j]作用于kv->v_cache[i][j]，获取最终的ans向量。
        // head_ans+=logits[j]×kv->v_cache[i][j]
        for (int j = 0; j < MAX_SEQ_LEN; j++)
        {
            vector_add(head_ans, &(kv->v_cache[i][j]), logits[j]);
        }
    }
}

3.2 paged_attention_v1_kernel

将上述的 C 串行代码改成cuda并行代码，并且充分发挥出GPU硬件算力，这并不是一件容易的事情。PagedAttention提供了两个version的kernal代码，我们先来看paged_attention_v1_kernel。

3.2.1 Input

const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]
const scalar_t* __restrict__ k_cache,   // [num_blocks, num_kv_heads, head_size/x, block_size, x] 
const scalar_t* __restrict__ v_cache,   // [num_blocks, num_kv_heads, head_size, block_size]

任何一个q k v向量的长度都是head_size。
q 存储没什么好说的，其实就是num_seqs×num_heads个向量存储在连续的内存空间。
k_cache中的 k 向量并没有存储在连续内存空间。其中，最后1维是x ，占用16字节（如果参数为FP16，则x=16/2；如果参数为FP32，则x=16/4）。
- block_size默认为16，也就是说，一个token block中16个token对应的16个 k 向量segment（16字节）存储在一块。
- 这儿太绕了，还是看图吧。
k_cache中的 v 向量也没有存储在连续内存空间。
- block_size默认为16，也就是说，一个token block中16个token对应的16个 k 向量以interleave的方式存储在一块。
- 还是看图吧。一图胜千言。

为了突出重点，图中num_head=1, block_size=4，只有两个token block：

在这里插入图片描述

3.2.2 workload partition

attention算子对并行计算很友好。我们来看一下paged_attention_v1_kernel中的并行计算任务划分，其实就是搞清楚cuda block、warp和thread各自承担了那些计算任务。

3.2.2.1 cuda grid

dim3 grid(num_heads, num_seqs, 1);

不同sequence中的attention任务互不干扰，同一sequence中不同head中的attention任务也互不干扰。

因此，paged_attention_v1_kernel将一个head中的attention任务交给一个cuda block完成。这样的话，不涉及不同cuda block间的同步，通信开销会比较小。

3.2.2.2 cuda block

1个cuda block 负责1个 head 中的attention任务。在PagedAttention中1个cuda block 中的任务划分比较复杂，不仅包含warp和thread，还包含group。整个过程可以分成两个主要阶段：1. 计算logits参数；2. 将logits参数作用于v_cache，得出结果。

cuda block中创建了128个cuda thread，也就是4个warp，1个warp中包含32个cuda thread。

阶段1，计算logits参数
- 单独拿出C串行伪代码：
```
// head_ans+=logits[j]×kv->v_cache[i][j]
for (int j = 0; j < MAX_SEQ_LEN; j++)
{
    logits[j] = vector_dot(&(q->vectors[i]), &(kv->k_cache[i][j]));
}
```
- 在paged_attention_v1_kernel中的任务划分：
  - warp
    
    一个或多个token block分配给一个warp执行。
    
    warp 会对 token block 按 x × block_size 字节的连续存储块依次进行处理。
  - thread group
    
    一个warp会划分成block_size个thread group，也就是说一个thread group负责处理一个token。
    
    thread group 会对 token 按 x 字节的连续存储块依次进行处理。
  - thread
    
    将token 中 x 字节的连续存储块划分成 VEC_SIZE = x / sizeof(type) / THREAD_GROUP_SIZE 的 VEC 进行处理。
    
    每一个 thread 负责 head_size / VEC_SIZE 个 VEC。
- 任务划分示意图
阶段2，将logits参数作用于v_cache，得出结果
- 单独拿出C串行伪代码：
```
for (int j = 0; j < MAX_SEQ_LEN; j++)
{
    vector_add(head_ans, &(kv->v_cache[i][j]), logits[j]);
}
```
- 在paged_attention_v1_kernel中的任务划分：
  - warp
    
    一个或多个token block分配给一个warp执行。
    
    warp 会对 token block 按 sizeof(type) × block_size 字节的连续存储块依次进行处理。
  - thread
    
    如果sizeof(type) × block_size 字节的连续存储块大于16字节，则将其切分为16字节的连续存储块。然后，将16字节的连续存储块依次分配给warp内线程进行处理。
    
    如果sizeof(type) × block_size 字节的连续存储块小于等于16字节，则将其依次分配给warp内线程进行处理。
- 任务划分示意图

3.2.3 伪代码

简单写一下paged_attention_v1_kernel的伪代码：

__device__ void paged_attention_kernel(
  scalar_t* __restrict__ out,             // [num_seqs, num_heads, max_num_partitions, head_size]
  const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]
  const scalar_t* __restrict__ k_cache,   // [num_blocks, num_kv_heads, head_size/x, block_size, x]
  const scalar_t* __restrict__ v_cache,   // [num_blocks, num_kv_heads, head_size, block_size]
  const int num_kv_heads,                 // [num_heads]
){
    __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
    /* 
    *	load q_vecs
    */
    __syncthreads();
    // Memory planning.
    float* logits = reinterpret_cast<float*>(shared_mem);
	for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; block_idx += NUM_WARPS) {
        K_vec k_vecs[NUM_VECS_PER_THREAD];
        for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
            for (int j = 0; j < NUM_VECS_PER_THREAD; j++){
                /* 
                *	load k_vecs
                */
            }
        }
        /* 
        *	compute q k dot
        */
        float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(q_vecs[thread_group_offset], k_vecs);
        logits[token_idx - start_token_idx] = mask ? 0.f : qk;
    }
    __syncthreads();
    /* 
    *	Compute softmax.
    */
    __syncthreads();
    float accs[NUM_ROWS_PER_THREAD];
    for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; block_idx += NUM_WARPS) {
        for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
            accs[i] += dot(logits_vec, v_vec);
        }
    }
    __syncthreads();
    /* 
    *	Perform accs reduction across warps.
    */
}

3.3 paged_attention_v2_kernel

paged_attention_v1_kernel将一个head中的attention任务交给一个cuda block完成。如果sequence中context比较长，cuda block承担的任务量将会比较大，甚至有可能出现硬件资源（共享内存，寄存器）不够的情况。这时候，就需要将head中的attention任务拆分成多个cuda block。在vLLM中，每512个token创建一个cuda block。

其他计算逻辑与paged_attention_v1_kernel完全相同。