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$
    • UPMEM中$\lambda$

• 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$	64$\lambda$	5.87 × 10^-6 s	0	514$\lambda$

  • 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$	192$\lambda$	5.87 × 10^-6 s	3.05 × 10^-6 s	642$\lambda$

  • 如果可以将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