Quantitative Analysis: PIM Chip Demands for LLAMA-7B inference-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。