Quantitative Analysis: PIM Chip Demands for LLAMA-7B inference

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
  - 推理过程中，保存中间结果
Die
- DieRAM
  - 数据 buffer，Controller broadcast 数据时用到。
Bank
- MACs
  - Multiply-And-Accumulate
  - 用于GEMV and GEMM
- BRAM
  - cache
- BMEM
  - capacity: 32MB

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

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