阅读 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:

  • 阅读 DistServe 论文
  • 阅读 Orca 论文