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大小),默认为 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完全相同。

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/news/627463.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

MCU、MPU、SOC简介

文章目录 前言一、MCU二、MPU三、SOC总结 前言 随着处理器技术的不断发展&#xff0c;CPU(Central Processing Unit)的发展逐渐出现三种分支&#xff0c;分别是MCU(Micro Controller Unit&#xff0c;微控制器单元) 和MPU&#xff08;Micro Processor Unit&#xff0c;微处理器…

Angular系列教程之管道

文章目录 管道的基本概念使用内置管道创建自定义管道总结 在Angular中&#xff0c;管道&#xff08;Pipe&#xff09;是一个非常重要的概念。它们允许我们对数据进行转换、格式化和显示&#xff0c;并且可以轻松地在模板中使用。本篇文章将介绍Angular中的管道概念&#xff0c;…

SL3036国产新品 48V/60V电动车里程增程器供电芯片

随着电动车的普及&#xff0c;里程焦虑成为了很多电动车用户面临的问题。为了解决这个问题&#xff0c;SL3036国产新品应运而生&#xff0c;它是一款48V/60V电动车里程增程器供电芯片。这款芯片的出现&#xff0c;为电动车用户提供了更加可靠的续航里程&#xff0c;让他们在出行…

使用AI自动生成PPT提高制作效率

使用AI自动生成PPT提高制作效率 在制作PPT方面&#xff0c;很多制作者都会轻易跳进一个怪圈&#xff1a;“我要制作一个关于关爱老人的PPT&#xff0c;该怎么做呢&#xff0c;有模板没有?”这个会涉及很多逻辑需要经过不断的思考&#xff0c;制作PPT要通过很多素材、使用技巧、…

3DMax的位图是什么? 3DMax的位图介绍

在3dmax建模中&#xff0c;使用贴图时的位图的频率是很高的。主要原因便就是位图就是我们平常说的图片&#xff0c;有各种格式的图片&#xff0c;能把这张图片贴到物体的表面&#xff0c;呈现效果进行渲染。 3damx的位图支持多种格式&#xff0c;比如jpg、png等等。 当然常用的…

el-table右固定最后一列显示不全或者是倒数第二列无边框线

问题图片&#xff1a; 解决方式1&#xff1a; >>>.el-table__row td:not(.is-hidden):last-child { border-left:1px solid #EBEEF5; } >>>.el-table__header th:not(.is-hidden):last-child{ border-left:1px solid #EBEEF5; } >>>.el-table__head…

C语言--质数算法和最大公约数算法

文章目录 1.在C语言中&#xff0c;判断质数的常见算法有以下几种&#xff1a;1.1.试除法&#xff08;暴力算法&#xff09;&#xff1a;1.2.优化试除法&#xff1a;1.3.埃拉托色尼筛法&#xff1a;1.4.米勒-拉宾素性检验&#xff1a;1.5.线性筛法&#xff1a;1.6.费马小定理&am…

最新可用GPT-3.5、GPT-4、Midjourney绘画、DALL-E3文生图模型教程【宝藏级收藏】

一、前言 ChatGPT3.5、GPT4.0、GPT语音对话、Midjourney绘画&#xff0c;文档对话总结DALL-E3文生图&#xff0c;相信对大家应该不感到陌生吧&#xff1f;简单来说&#xff0c;GPT-4技术比之前的GPT-3.5相对来说更加智能&#xff0c;会根据用户的要求生成多种内容甚至也可以和…

宋仕强论道之华强北电子交易中心(四十三)

最近中国电子CECC牵头的电子交易中心在前海挂牌&#xff0c;这给华强北带来了一些威胁&#xff0c;也蚕食了部分华强北市场。事实上&#xff0c;电子元器件市场是一个非常巨大且复杂的市场&#xff0c;包括了各种各样的产品和配套的服务。比如说存储产品这一块&#xff0c;标准…

在Overleaf中解决IEEE LaTeX模板不能显示中文问题

解决IEEE的Latex模板不能显示中文的问题 写在最前面编译器选择XeLatex导入CTeX包IEEE单栏转换为双栏如何在Overleaf中解决IEEE LaTeX模板显示中文问题&#xff1a;一些其他的补充引言问题描述准备工作为什么中文字符在IEEE LaTeX模板中显示有问题——了解LaTeX编码的基础概念 关…

写点东西《最佳 Web 框架不存在 》

写点东西《&#x1f947;最佳 Web 框架不存在 &#x1f6ab;》 TLDR&#xff1b;您选择的 Web 应用程序框架并不重要。嗯&#xff0c;它很重要&#xff0c;但并不像其他人希望您相信的那样重要。 2024 年存在如此多的库和框架&#xff0c;而且最好的库和框架仍然备受争议&…

C语言——小细节和小知识10

一、全局变量和局部变量 1、引例 当全局变量和局部变量名字相同的情况下&#xff0c;局部变量优先。 #include <stdio.h>int num 10;int main() {int num 0;printf("%d\n", num);return 0; } 运行结果 2、介绍 在C语言中&#xff0c;当局部变量和全局变…

高光谱分类论文解读分享之基于生成对抗性少数过采样的高光谱图像分类

IEEE TGRS 2022&#xff1a;基于生成对抗性少数过采样的高光谱图像分类 题目 Generative Adversarial Minority Oversampling for Spectral–Spatial Hyperspectral Image Classification 作者 Swalpa Kumar Roy , Student Member, IEEE, Juan M. Haut , Senior Member, IE…

【野火i.MX6NULL开发板】ARM-GCC 和开发板的 HelloWorld(ubuntu主机和野火开发板debian交叉编译)、开发板的/mnt里没文件

0、前言 参考资料&#xff1a; 《野火 Linux 基础与应用开发实战指南基于 i.MX6ULL 系列》PDF 第24章 参考视频&#xff1a; https://www.bilibili.com/video/BV1JK4y1t7io?p26&vd_sourcefb8dcae0aee3f1aab700c21099045395 注意&#xff0c;一定要记得把虚拟机的网络适配…

J3-DenseNet实战

&#x1f368; 本文为&#x1f517;365天深度学习训练营 中的学习记录博客&#x1f356; 原作者&#xff1a;K同学啊 | 接辅导、项目定制 目录 环境步骤环境设置数据准备图像信息查看 模型构建模型训练模型效果展示 总结与心得体会 环境 系统: Linux语言: Python3.8.10深度学习…

接口自动化神器 apin【 快速入门篇】

关于自动化测试&#xff0c;这些年经历了太多的坑&#xff0c;有被动的坑&#xff0c;也有自己主动挖的坑&#xff0c;在这里做了一些总结。 一、apin 介绍及安装 1、什么是 apin apin是一个无需写代码&#xff0c;就可以进行接口自动化测试的框架&#xff0c;只需要通过jso…

【线路图】世微AP5160宽电压降压型恒流芯片 LED电源 带调光SOT23-6

这是一款14-18V 3A 电流的PCB设计方案. 运用的是世微AP5160 电源驱动IC,这是一款效率高&#xff0c;稳定可靠的 LED 灯恒流驱动控制芯片&#xff0c;内置高精度比较器&#xff0c;固定 关断时间控制电路&#xff0c;恒流驱动电路等&#xff0c;特别适合大功率 LED 恒流驱动。 …

59.说一下 spring 的事务隔离?

spring 的事务隔离有什么作用? 用来解决并发事务所产生一些问题,并发会产生什么问题? 1.脏读2.不可重复度3.幻影读事务隔离的概念 通过设置隔离级别可解决在并发过程中产生的那些问题分别举例说明 1.脏读 上述图表示:一个事务,读取了另一个事务中没有提交的数据,会在…

【AD-3D预览-颜色更换和模型操作】AD打开3D预览模式,PCB板子显示蓝色,如何更改为绿色

问题&#xff1a;AD软件的3D预览模式中&#xff0c;PCB的颜色和其他不一样&#xff0c;显示的是蓝色&#xff0c;而且正面可以看到走线&#xff0c;背面看不到 原因&#xff1a; 这是因为选择的查看模式不一样所导致的。 在这个位置可以查看当前所使用的模式是什么&#xff0…

IPKISS ------ 远程服务器 IPKISS 内置示例安装问题

IPKISS ------ 远程服务器示例安装问题 引言正文 引言 很多时候&#xff0c;如果我们在服务器上使用管理员权限安装了 IPKISS 证书&#xff0c;而我们使用个人账号登录服务器时有时候会显示如下界面&#xff1a; 我们会看到这个 PyCharm (Luceda Academy) 是灰色的。那么该怎…