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大小），默认为 16self.token_ids = [_BLANK_TOKEN_ID] * block_size	// 初始化为TOKEN_ID为-1，表示无效tokenself.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_tokensself.token_ids[curr_idx:curr_idx + len(token_ids)] = token_idsself.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 CPUself.block_number = block_number	// PhysicalTokenBlock id, 由零递增self.block_size = block_size	// block大小，与LogicalTokenBlock中block_size相等，默认为 16self.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 whenrequested. When a block is freed, its reference count is decremented. Ifthe 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 CPUself.block_size = block_size	// block大小，默认为 16self.num_blocks = num_blocks	// # Initialize the free blocks.self.free_blocks: BlockTable = []	// 一个链表，存储空闲的PhysicalTokenBlockfor 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_sizeself.num_total_gpu_blocks = num_gpu_blocksself.num_total_cpu_blocks = num_cpu_blocksself.watermark = watermarkassert watermark >= 0.0self.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 theblock 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_idself.prompt = promptself.block_size = block_sizeself.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语言串行代码非常简单，容易理解。

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

   计算出来MAX_SEQ个logit，代表了 q 向量和 k 向量之间的相关性。

2. 对MAX_SEQ个logit进行softmax

3. 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进行softmaxsoft_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完全相同。