1 Architecture
如果将LLAMA-7B模型参数量化为4bit,则存储模型参数需要3.3GB。那么,至少PIM chip 的存储至少要4GB。
- AiM单个bank为32MB,单个die 512MB,至少需要8个die的芯片。
- 8个die集成在一个芯片上。
- 提供8×16bank级别的访存带宽。
- 整个推理过程完全下放至PIM。
- CPU把 prompt 传给 Controller
- Controller 控制推理过程,将推理出的token返回给CPU
- Controller
- ALUs
- 处理softmax、Norm和向量乘等。
- CRAM
- cache
- CMEM
- 推理过程中,保存中间结果
- ALUs
- Die
- DieRAM
- 数据 buffer,Controller broadcast 数据时用到。
- DieRAM
- Bank
- MACs
- Multiply-And-Accumulate
- 用于GEMV and GEMM
- BRAM
- cache
- BMEM
- capacity: 32MB
- MACs
2 Data partition
2.1 LLAMA-7B
2.2 Model parameter
在 batch_size=1 的情况下
- prefill 阶段,嵌入prompt,此时为GEMM算子
- decode 阶段,推理出一个个token,此时为GEMV算子
模型参数划分就是将上图中的矩阵划分至8×16个bank中。
2.2.1 一维划分
2.2.2 二维划分
- 在分布式计算场景下的常用划分
- 优势:通信量小。但是,要求计算节点间存在通信能力。
- 在PIM场景下,无法假设bank间存在通信能力,此时,Controller的reduction开销会很大。
2.3 KV cache
2.3.1 attention
- NUM_HEAD 个相互独立的 attention 操作
- max sequence length = (8×16×32×1024×1024 - 6607077376/2)/(32×4096/2) = 15128
2.3.2 a bank for a head
- 一个head attention由一个bank执行
- 优势:Controller 与 banks 通信少
- 劣势:
- NUM_HEAD < NUM_BANK,3/4的bank访存带宽和算力被浪费。
- bank内不仅要支持MAC,还要支持softmax。
2.3.3 multiple banks for a head
- 一个head attention由多个bank执行
- 优势:所有的bank访存带宽和算力得到利用。
- 劣势:
- Controller 与 banks 通信开销变大
- Controller 需要进行softmax和reduction。
3 Demands
本节使用量化方法来分析PIM chip,希望能够回答以下几个问题:
CRAM、CMEM、BRAM做多大合适?
Bank级并行带宽需要多少?并行算力需要多少?
Controller 如何与bank通信?通信带宽需要多少?
Controller 需要提供多大的算力?
PIM chip面向端侧推理,一般来说,推理的batch size = 1。LLM在推理时,可以大致分为两个阶段:
- prefill
- 在prefill阶段,模型嵌入prompt。
- 假设嵌入prompt的长度为N,则在这个过程中模型参数会被复用N次。
- 典型算子为GEMM
- decode
- 在decode阶段,模型以自回归的方式推理出一个个token
- 假设推理出S个token
- 每推理出一个token,则在这个过程中都必须扫描一遍模型参数和kv cache。
- 典型算子为GEMV
3.1 bandwidth
显然,LLM模型在decode阶段的瓶颈在于访存带宽。定量分析decode过程,也就可以分析出在给定访存带宽下,模型推理的速度。
3.1.1 hypothesis
假设推理的sequence的长度为L
在decode阶段,Controller和bank内的算力均可以吃满访存带宽。
- 对于GEMV算子,Operational intensity = 2 Ops/weight byte
- 这个假设完全合理。
并行intra-bank bandwidth 总带宽为BM
- 对于AiM,BM = 8 × 512 GB/s
Controler-bank bandwidth 总带宽为CBM
一般来说,CBM << BM
对于AiM,CBM = 8 × 32 GB/s
这儿应该进行更加精细的讨论
- 每次 Controler-bank 通信的基础开销(时延)设置为 λ \lambda λ
intra-Controler bandwidth 带宽为CM
- 能够达到类似CPU 100GB/s的访存带宽?
在decode过程,推理出一个token的时延 = bank内并行访存(GEMV)的时延 + Controler-bank 通信的时延 + Controler 内访存(softmax、Norm 和 reduction 等)的时延
3.1.2 intra-bank bandwidth
bank内并行访存(GEMV)的时延包含两部分:模型参数相关的GEMV的时延和kv cache相关的GEMV的时延。
模型参数相关的GEMV的时延
- 这个时延非常好算,其实就是 模型总参数量/BM。
- 对于AiM,这个时延为 (6607077376 / 2 / 1024 / 1024 / 1024)/(BM) = 7.5 × 10^-4 s
kv cache相关的GEMV的时延
- kv cache 大小
- L×H×NUM_LAYER×2
- 与kv cache相关的attention算子的时延计算就比较复杂,因为其有两者计算方案。
- a bank for a head
- NUM_HEAD 的数量为32,bank数量为8×16
- bank并行访存带宽的利用率为25%
- 时延 = ((L × H + NUM_HEAD × L + L × H) × NUM_LAYER / 2 / 1024 / 1024 / 1024)/(BM / 4)
- 如果L=4096,时延 = 4.9 × 10^-4 s
- multiple banks for a head
- 4 banks for a head
- 时延 = ((L × H + L × H) × NUM_LAYER / 2 / 1024 / 1024 / 1024)/(8 × 512)
- 如果L=4096,时延 = 1.2 × 10^-4 s
- kv cache 大小
3.1.3 Controler-bank bandwidth
模型参数相关的通信时延
- 通信次数
- 运算一个GEMV算子,需要两次通信,向bank发送向量,取回结果。
- 通信次数 2 × (1 + NUM_LAYER × 7)
- 通信量
- (V + H) + NUM_LAYER × (4 × 2 × H + 3 × (H + DIM_MLP))
- 通信时延
- (2 × (1 + NUM_LAYER × 7)) λ \lambda λ + ((V + H) + NUM_LAYER × (4 × 2 × H + 3 × (H + DIM_MLP))) /2 / 1024 / 1024 / 1024 / CBM
- UPMEM中 λ \lambda λ=0.0001s,AiM中CBM=8 × 32 GB/s,此时,时延 = 450 λ \lambda λ + 4.61×10^-6 s
- 通信次数
kv cache相关的通信时延
a bank for a head
通信次数
2 × NUM_LAYER
通信量
NUM_LAYER × (2H)
通信时延
(2 × NUM_LAYER) λ \lambda λ + NUM_LAYER × (2H) /2 / 1024 / 1024 / 1024 / (CBM/4)
64 λ \lambda λ + 1.91 × 10^-6 s
multiple (4) banks for a head
通信次数
2 × NUM_LAYER × 2
通信量
NUM_LAYER × (4H + NUM_HEAD × L + NUM_HEAD × L + 4H)
通信时延
(2 × NUM_LAYER × 2) λ \lambda λ + NUM_LAYER × (4H + NUM_HEAD × L + NUM_HEAD × L + 4H) /2 / 1024 / 1024 / 1024 / CBM
192 λ \lambda λ + 1.72 × 10^-5 s
3.1.4 intra-Controler bandwidth
模型参数相关的GEMV的时延
- 主要是做处理softmax、Norm和向量乘等。
- 处理数据量
- V + NUM_LAYER × (H + H + H + 2 × DIM_MLP + H)
- 处理时延
- (V + NUM_LAYER × (H + H + H + 2 × DIM_MLP + H)) /2 / 1024 / 1024 / 1024 / CM
- 5.87 × 10^-6 s
kv cache相关的GEMV的时延
a bank for a head
- 无
multiple (4) banks for a head
softmax 和 reduction
处理数据量
NUM_LAYER × (L + 4 × H)
处理时延
NUM_LAYER × (L + 4 × H) /2 / 1024 / 1024 / 1024 / CM
3.05 × 10^-6 s
3.1.5 summary
sequence len = 4096,推理出一个token的总时延
a bank for a head
intra-bank 模型参数 intra-bank qkv Controler-bank 模型参数 Controler-bank qkv intra-Controler 模型参数 intra-Controler qkv total 7.5 × 10^-4 s 4.9 × 10^-4 s 450 λ \lambda λ + 4.61×10^-6 s 64 λ \lambda λ + 1.91 × 10^-6 s 5.87 × 10^-6 s 0 514 λ \lambda λ+1.25 × 10^-3 s multiple (4) banks for a head
intra-bank 模型参数 intra-bank qkv Controler-bank 模型参数 Controler-bank qkv intra-Controler 模型参数 intra-Controler qkv total 7.5 × 10^-4 s 1.2 × 10^-4 s 450 λ \lambda λ + 4.61×10^-6 s 192 λ \lambda λ + 1.72 × 10^-5 s 5.87 × 10^-6 s 3.05 × 10^-6 s 642 λ \lambda λ+1.27 × 10^-3 s 如果可以将Controler-bank通信时延优化掉,1s 可以推理出500+ token。这时候,性能是冗余的。
Controler要放在PIM chip上,否则Controler-bank通信基础开销会成为系统瓶颈。
削减成本,将上述三种带宽均减小一个数量级,系统吞吐 50+ token/s,性能也能满足需求。
3.2 Computing
LLM模型在prefill阶段的瓶颈在于硬件算力。
3.2.1 Controler
- Controler 部分承担softmax、Norm 和 reduction 访存密集性算子
- Controler 的算力只要能吃满 Controler 的访存带宽就可以。
- CRAM容量
- L最长16K,能够放入L长度的向量进行softmax就可以
- CRAM容量最小 8KB,设置64KB是合理的。
- 当然,如果追求更大的并行度,设置 NUM_HEAD × 8KB = 256 KB
- CMEM容量
- 在prefill阶段,要能够放入L×L的矩阵,最小128MB
- 设置256MB是合理的
3.2.2 bank
bank 承担GEMM算子
嵌入prompt的长度为N,模型参数加载进入cache可以重用N次
核心还是MAC
可以将 N 以 batch_size (例如16) 为粒度进行切分,以tile的方式进行GEMM
- 可以设置BRAM的大小为 16 KB ~ 64 KB,更大的BRAM可以允许更大的tile。
需要算力:BM×2×batch_size ops
3.3 conclusion
如果decode阶段需要50+ token/s的推理速度:
CRAM、CMEM、BRAM做多大合适?
CRAM CMEM BRAM 256 KB 256MB 16 KB ~ 64 KB Bank级并行带宽需要多少?并行算力需要多少?
512 GB/s
512×2 ~ 512×2×16 Gops
prefill 阶段的embedding速度和并行算力相关。
512×2 Gops 对应50+ token/s的embedding速度
512×2×16 Gops 对应50×16+ token/s的embedding速度
Controller 如何与bank通信?通信带宽需要多少?
- Controller 必须在片内,降低通信时延
- 4 ~ 32GB/s
Controller 需要提供多大的访存带宽和算力?
- 16 GB/s
- 32 Gops