阅读 Paged Attention 论文
Q, K, V都是词向量通过W线性变换得到的
当前Token的Q和所有的K作点积
K是Token的关键词
V是Token的语义信息
最后通过softmax得到每个词的概率,选择下一个Token
Q和All Keys进行查询,然后和All Values进行计算
只有Q是当前的,K和V都是过去的所有的,可以复用。
避免内外碎片
为何需要共享
-
Parallel Samling(并行采样) 对同一个Prompt,输出不同结果
-
Beam Search 高级的解码机制 在每一步保留多个得分最好的(多个)候选序列,逐步扩展这些序列,以找到整体概率最高的输出。
Paged Attention
使用分页机制,Block as pages, tokens as bytes, requests as processes.
调度设计
Preemptive request scheduling 驱逐出一些请求(
Background
Prefill阶段和Decode阶段
Prefill计算密集型
Decode阶段线性,每次生成一个Token
结束触发条件:EOS/最大长度
多个请求进行Batch处理
- 用户请求并非同时到达
- 不同Prompt生成长度区别很大
- 需要Padding
Cellular Batching Iteration Level
已完成的移除,新的加入进来
Memory Challenges
对于13B模型一个Token的KV Cache就有800KB 如果一个请求2048token,1.6GB!
GPU计算能力增长比内存快
复杂的 Decoding Algorithms
共享内存,减小开销
对未知的Input和Output进行Scheduling
Method
集中调度器,协调多个Worker执行。
允许连续的KV储存在非连续的空间。
16 token的Block Size GPU Kernel 逐个处理 Query
基于Block去算,不是对单个Token去算。 按块进行处理
KV Cache Manager 物理内存空间无需提前预留 通过 PagedAttention 加 KV Cache
还可以和CPU Mem进行换页!
Block Table
Physical Block Number -> filled数量
首先从活跃的请求中选择一批进行处理 将Prompt中所有Token进行拼接然后处理?
Block Size的选择。
Decoding
- Parallel Sampling
生成多个候选输出 使用CoW机制
- Beam Search
Prompt和生成都能共享
- Shared Prefix(System Prompt)
调度和抢占
FCFS(First Come First Serve)
驱逐哪个Block,如何恢复?
一个请求可能使用到多个Block All-or-Nothing 驱逐
两种可行策略
-
Swapping 拷贝到CPU内存
-
Recomputation? 没看懂
4.6 分布式执行
Megatron-LM Like
所有Worker共享映射表
每个Worker只存一小部分KV Cache
Scheduler准备控制信息, 请求的tokenID和Block Table。 广播给Worker
all-reduce
操作?无需Scheeduler
Implementation
FastAPI,OpenAI interface
- Scheduler / Block Manager (Python)
- PagedAttention(CUDA)
CUDA Kernel 是运行在GPU上的并行函数
Kernel Level Optimization
- Fused reshap and block write
- Fusing block read and attention
- Fused block copy
- CoW Mechanism
fork
append
free
支持 Decoding 算法
Evaluation
Experimental Setup
Baseline: FasterTransformer 设置最大的Batch Size
Orca(Oracle): 知道输出长度 Orca(Pow2): 预留2^n长度 Orca(Max): 预留2048
Parallel Sampling and Beam Search
Chatbot,带上历史记录
消融实验(Ablation Studies)
把模型或系统里的某个组件(或特性)“拿掉/替换/简化”之后,重新做实验,比较性能变化,从而证明该组件的重要性。
Kernel还是有延迟 Swap对于小Block Size来讲,PCIe通信开销太大。
Enjoy Reading This Article?
Here are some more articles you might like to read next: