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: 32MB2 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) = 151282.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。 典型算子为GEMV3.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/sControler-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 skv 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 s3.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 skv 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 skv 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 smultiple (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 KBBank级并行带宽需要多少？并行算力需要多少？
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/sController 需要提供多大的访存带宽和算力？
16 GB/s 32 Gops