视频网站后台功能,可以安装两个wordpress,盐城市城南新区建设局网站,办个网站需要多少钱1 Overview
PagedAttention灵感来自于操作系统中虚拟内存和分页的经典思想#xff0c;它可以允许在非连续空间立存储连续的KV张量。具体来说#xff0c;PagedAttention把每个序列的KV缓存进行了分块#xff0c;每个块包含固定长度的token#xff0c;而在计算attention时可…1 Overview
PagedAttention灵感来自于操作系统中虚拟内存和分页的经典思想它可以允许在非连续空间立存储连续的KV张量。具体来说PagedAttention把每个序列的KV缓存进行了分块每个块包含固定长度的token而在计算attention时可以高效地找到并获取那些块。 2 Block management
相比于RaggedAttentionPagedAttention其实就是维护了一个逻辑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只能从左到右appendassert 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(devicedevice,block_numberi,block_sizeblock_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的BlockTableBlockTable 维护者 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 attentionNUM_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进行softmaxsoft_max(logits);head_ans (ans-vectors[i]);// logits[j]作用于kv-v_cache[i][j]获取最终的ans向量。// head_anslogits[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则x16/2如果参数为FP32则x16/4。 block_size默认为16也就是说一个token block中16个token对应的16个 k 向量segment16字节存储在一块。这儿太绕了还是看图吧。 k_cache中的 v 向量也没有存储在连续内存空间。 block_size默认为16也就是说一个token block中16个token对应的16个 k 向量以interleave的方式存储在一块。还是看图吧。一图胜千言。
为了突出重点图中num_head1, block_size4只有两个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个warp1个warp中包含32个cuda thread。 阶段1计算logits参数 单独拿出C串行伪代码 // head_anslogits[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_castfloat*(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_dotscalar_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完全相同。
